import { useState } from "react";
import { useMutation } from "@tanstack/react-query";
import { Button } from "@/components/ui/button";
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
import { Sparkles, Loader2, PenLine, CheckCircle, BarChart3, ChevronDown } from "lucide-react";
import { coachingService } from "@/services/coaching.service";
import { useToast } from "@/hooks/use-toast";

type Mode = "improve" | "grammar" | "band" | null;

interface Props {
  text: string;
  task_type?: string;
}

export default function AiWritingHelper({ text, task_type = "ielts_writing" }: Props) {
  const [open, setOpen] = useState(false);
  const [activeMode, setActiveMode] = useState<Mode>(null);
  const [showResult, setShowResult] = useState(false);
  const { toast } = useToast();

  const mutation = useMutation({
    mutationFn: (mode: NonNullable<Mode>) =>
      coachingService.writingHelp({
        text: text.trim(),
        task_type: `${task_type}:${mode}`,
      }),
    onSuccess: () => setShowResult(true),
    onError: (err: Error) => {
      toast({
        title: "Writing help failed",
        description: err.message || "Could not analyze your writing. Try again.",
        variant: "destructive",
      });
    },
  });

  const handleAction = (mode: Mode) => {
    if (!mode) return;
    if (!text.trim()) {
      toast({
        title: "Add some text first",
        description: "Enter your draft in the text area so AI can analyze it.",
        variant: "destructive",
      });
      return;
    }
    setActiveMode(mode);
    setShowResult(false);
    mutation.mutate(mode);
  };

  const loading = mutation.isPending;

  return (
    <Collapsible open={open} onOpenChange={setOpen}>
      <CollapsibleTrigger asChild>
        <Button variant="outline" size="sm" className="w-full justify-between mt-3">
          <span className="flex items-center gap-2">
            <Sparkles className="h-3.5 w-3.5 text-primary" />
            AI Writing Helper
          </span>
          <ChevronDown className={`h-4 w-4 transition-transform ${open ? "rotate-180" : ""}`} />
        </Button>
      </CollapsibleTrigger>
      <CollapsibleContent className="mt-3 space-y-3">
        <div className="flex gap-2">
          <Button variant="outline" size="sm" onClick={() => handleAction("improve")} disabled={loading}>
            <PenLine className="h-3.5 w-3.5 mr-1" /> Improve my draft
          </Button>
          <Button variant="outline" size="sm" onClick={() => handleAction("grammar")} disabled={loading}>
            <CheckCircle className="h-3.5 w-3.5 mr-1" /> Check grammar
          </Button>
          <Button variant="outline" size="sm" onClick={() => handleAction("band")} disabled={loading}>
            <BarChart3 className="h-3.5 w-3.5 mr-1" /> Estimate band score
          </Button>
        </div>

        {loading && (
          <div className="flex items-center gap-2 text-sm text-muted-foreground rounded-lg bg-muted/50 p-3">
            <Loader2 className="h-4 w-4 animate-spin text-primary" /> AI is analyzing your writing...
          </div>
        )}

        {showResult && !loading && mutation.data && activeMode === "improve" && (
          <div className="space-y-3">
            {mutation.data.feedback && (
              <div className="rounded-lg border bg-muted/30 p-3">
                <p className="text-xs font-semibold text-primary mb-1 flex items-center gap-1">
                  <Sparkles className="h-3 w-3" /> Feedback
                </p>
                <p className="text-sm text-muted-foreground">{mutation.data.feedback}</p>
              </div>
            )}
            {mutation.data.improved && (
              <div className="rounded-lg border bg-muted/30 p-3">
                <p className="text-xs font-semibold text-primary mb-1 flex items-center gap-1">
                  <Sparkles className="h-3 w-3" /> Improved Version
                </p>
                <p className="text-sm">{mutation.data.improved}</p>
              </div>
            )}
          </div>
        )}

        {showResult && !loading && mutation.data && activeMode === "grammar" && (
          <div className="rounded-lg border bg-muted/30 p-3 space-y-2">
            <p className="text-xs font-semibold text-primary mb-1 flex items-center gap-1">
              <Sparkles className="h-3 w-3" /> Grammar notes
            </p>
            {(mutation.data.grammar_notes?.length ?? 0) > 0 ? (
              mutation.data.grammar_notes!.map((note, i) => (
                <div key={i} className="text-sm border-l-2 border-warning pl-2">
                  <p className="text-muted-foreground">{note}</p>
                </div>
              ))
            ) : (
              <p className="text-sm text-muted-foreground">No grammar issues flagged.</p>
            )}
            {mutation.data.feedback ? (
              <p className="text-xs text-muted-foreground pt-2 border-t">{mutation.data.feedback}</p>
            ) : null}
          </div>
        )}

        {showResult && !loading && mutation.data && activeMode === "band" && (
          <div className="rounded-lg border bg-muted/30 p-3">
            <p className="text-xs font-semibold text-primary mb-1 flex items-center gap-1">
              <Sparkles className="h-3 w-3" /> Estimated band / assessment
            </p>
            <p className="text-sm text-muted-foreground">{mutation.data.feedback}</p>
            {mutation.data.improved ? (
              <p className="text-sm mt-2 pt-2 border-t">{mutation.data.improved}</p>
            ) : null}
          </div>
        )}
      </CollapsibleContent>
    </Collapsible>
  );
}