import { useState } from "react";
import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { Badge } from "@/components/ui/badge";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { Textarea } from "@/components/ui/textarea";
import { 
  Brain, 
  Sparkles, 
  FileText, 
  Search, 
  Loader2, 
  TrendingUp,
  Lightbulb,
  Target,
  CheckCircle,
  AlertCircle
} from "lucide-react";

interface AIAssistantProps {
  onDocumentSelect?: (documentId: number) => void;
}

interface EnhancedSearchResult {
  results: any[];
  enhancedQuery?: {
    enhancedQuery: string;
    intent: string;
    keywords: string[];
    suggestions: string[];
  };
  searchInsights?: {
    totalResults: number;
    avgRelevanceScore: number;
    modalResultsCount: number;
    localResultsCount: number;
  };
}

interface ResearchSynthesis {
  synthesis: string;
  keyFindings: string[];
  gaps: string[];
  recommendations: string[];
}

export default function AIAssistant({ onDocumentSelect }: AIAssistantProps) {
  const [query, setQuery] = useState("");
  const [selectedDocuments, setSelectedDocuments] = useState<number[]>([]);
  const [analysisText, setAnalysisText] = useState("");
  const queryClient = useQueryClient();

  // Enhanced AI Search
  const aiSearchMutation = useMutation({
    mutationFn: async (searchQuery: string): Promise<EnhancedSearchResult> => {
      const response = await fetch("/api/ai-search", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ 
          query: searchQuery, 
          maxResults: 10,
          useQueryEnhancement: true 
        }),
      });
      if (!response.ok) throw new Error("Enhanced search failed");
      return response.json();
    },
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: ["/api/search"] });
    },
  });

  // Query Enhancement
  const queryEnhancementMutation = useMutation({
    mutationFn: async (originalQuery: string) => {
      const response = await fetch("/api/enhance-query", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ query: originalQuery }),
      });
      if (!response.ok) throw new Error("Query enhancement failed");
      return response.json();
    },
  });

  // Document Analysis
  const documentAnalysisMutation = useMutation({
    mutationFn: async ({ content, analysisType }: { content: string; analysisType: string }) => {
      const response = await fetch("/api/analyze-document", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ content, analysisType }),
      });
      if (!response.ok) throw new Error("Document analysis failed");
      return response.json();
    },
  });

  // Research Synthesis
  const researchSynthesisMutation = useMutation({
    mutationFn: async ({ query, documentIds }: { query: string; documentIds: number[] }): Promise<ResearchSynthesis> => {
      const response = await fetch("/api/research-synthesis", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ query, documentIds }),
      });
      if (!response.ok) throw new Error("Research synthesis failed");
      return response.json();
    },
  });

  // Generate Embeddings
  const embeddingsMutation = useMutation({
    mutationFn: async (input: string) => {
      const response = await fetch("/api/embeddings", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ input }),
      });
      if (!response.ok) throw new Error("Embedding generation failed");
      return response.json();
    },
  });

  const handleEnhancedSearch = () => {
    if (!query.trim()) return;
    aiSearchMutation.mutate(query);
  };

  const handleQueryEnhancement = () => {
    if (!query.trim()) return;
    queryEnhancementMutation.mutate(query);
  };

  const handleDocumentAnalysis = (analysisType: string) => {
    if (!analysisText.trim()) return;
    documentAnalysisMutation.mutate({ content: analysisText, analysisType });
  };

  const handleResearchSynthesis = () => {
    if (!query.trim() || selectedDocuments.length === 0) return;
    researchSynthesisMutation.mutate({ query, documentIds: selectedDocuments });
  };

  const handleGenerateEmbeddings = () => {
    if (!query.trim()) return;
    embeddingsMutation.mutate(query);
  };

  return (
    <div className="space-y-6">
      <Card className="border-gradient-to-r from-blue-200 to-purple-200 dark:from-blue-800 dark:to-purple-800">
        <CardHeader>
          <CardTitle className="flex items-center gap-2 text-xl">
            <Brain className="w-6 h-6 text-blue-600" />
            AI Research Assistant
            <Badge variant="secondary" className="ml-2">Powered by Nebius & Modal</Badge>
          </CardTitle>
        </CardHeader>
        <CardContent>
          <Tabs defaultValue="search" className="w-full">
            <TabsList className="grid grid-cols-4 w-full mb-6">
              <TabsTrigger value="search" className="flex items-center gap-2">
                <Search className="w-4 h-4" />
                Smart Search
              </TabsTrigger>
              <TabsTrigger value="analysis" className="flex items-center gap-2">
                <FileText className="w-4 h-4" />
                Analysis
              </TabsTrigger>
              <TabsTrigger value="synthesis" className="flex items-center gap-2">
                <Lightbulb className="w-4 h-4" />
                Synthesis
              </TabsTrigger>
              <TabsTrigger value="embeddings" className="flex items-center gap-2">
                <Sparkles className="w-4 h-4" />
                Embeddings
              </TabsTrigger>
            </TabsList>

            {/* Enhanced Search Tab */}
            <TabsContent value="search" className="space-y-4">
              <div className="space-y-3">
                <div className="flex gap-2">
                  <Input
                    placeholder="Enter research query for AI-enhanced search..."
                    value={query}
                    onChange={(e) => setQuery(e.target.value)}
                    onKeyDown={(e) => e.key === "Enter" && handleEnhancedSearch()}
                    className="flex-1"
                  />
                  <Button 
                    onClick={handleEnhancedSearch}
                    disabled={!query.trim() || aiSearchMutation.isPending}
                  >
                    {aiSearchMutation.isPending ? (
                      <Loader2 className="w-4 h-4 animate-spin" />
                    ) : (
                      <Search className="w-4 h-4" />
                    )}
                  </Button>
                </div>

                <div className="flex gap-2">
                  <Button 
                    variant="outline" 
                    size="sm"
                    onClick={handleQueryEnhancement}
                    disabled={!query.trim() || queryEnhancementMutation.isPending}
                  >
                    {queryEnhancementMutation.isPending ? (
                      <Loader2 className="w-3 h-3 animate-spin mr-1" />
                    ) : (
                      <Target className="w-3 h-3 mr-1" />
                    )}
                    Enhance Query
                  </Button>
                </div>
              </div>

              {/* Query Enhancement Results */}
              {queryEnhancementMutation.data && (
                <Card className="bg-blue-50 dark:bg-blue-900/20 border-blue-200 dark:border-blue-800">
                  <CardContent className="pt-4">
                    <h4 className="font-semibold text-blue-900 dark:text-blue-100 mb-2">Enhanced Query</h4>
                    <p className="text-sm mb-3 font-mono bg-white dark:bg-gray-800 p-2 rounded">
                      {queryEnhancementMutation.data.enhancedQuery}
                    </p>
                    <div className="grid grid-cols-1 md:grid-cols-2 gap-3 text-xs">
                      <div>
                        <span className="font-medium text-blue-800 dark:text-blue-200">Intent:</span>
                        <span className="ml-2">{queryEnhancementMutation.data.intent}</span>
                      </div>
                      <div>
                        <span className="font-medium text-blue-800 dark:text-blue-200">Keywords:</span>
                        <div className="flex flex-wrap gap-1 mt-1">
                          {queryEnhancementMutation.data.keywords.map((keyword: string, i: number) => (
                            <Badge key={i} variant="outline" className="text-xs">
                              {keyword}
                            </Badge>
                          ))}
                        </div>
                      </div>
                    </div>
                  </CardContent>
                </Card>
              )}

              {/* Enhanced Search Results */}
              {aiSearchMutation.data && (
                <Card>
                  <CardHeader>
                    <CardTitle className="flex items-center gap-2 text-lg">
                      <TrendingUp className="w-5 h-5 text-green-600" />
                      AI-Enhanced Results
                      {aiSearchMutation.data.searchInsights && (
                        <Badge variant="secondary">
                          {aiSearchMutation.data.searchInsights.totalResults} results
                        </Badge>
                      )}
                    </CardTitle>
                  </CardHeader>
                  <CardContent className="space-y-3">
                    {aiSearchMutation.data.searchInsights && (
                      <div className="grid grid-cols-2 md:grid-cols-4 gap-3 p-3 bg-gray-50 dark:bg-gray-800 rounded-lg text-sm">
                        <div>
                          <span className="font-medium">Avg Relevance:</span>
                          <span className="ml-1 text-green-600">
                            {(aiSearchMutation.data.searchInsights.avgRelevanceScore * 100).toFixed(1)}%
                          </span>
                        </div>
                        <div>
                          <span className="font-medium">Modal Results:</span>
                          <span className="ml-1">{aiSearchMutation.data.searchInsights.modalResultsCount}</span>
                        </div>
                        <div>
                          <span className="font-medium">Local Results:</span>
                          <span className="ml-1">{aiSearchMutation.data.searchInsights.localResultsCount}</span>
                        </div>
                        <div>
                          <span className="font-medium">Total:</span>
                          <span className="ml-1">{aiSearchMutation.data.searchInsights.totalResults}</span>
                        </div>
                      </div>
                    )}

                    <div className="space-y-2 max-h-96 overflow-y-auto">
                      {aiSearchMutation.data.results.map((result: any, index: number) => (
                        <Card key={index} className="p-3 hover:bg-gray-50 dark:hover:bg-gray-800">
                          <div className="flex justify-between items-start mb-2">
                            <h5 className="font-medium text-sm">{result.title}</h5>
                            <div className="flex items-center gap-2">
                              {result.relevanceScore && (
                                <Badge variant="outline" className="text-xs">
                                  {(result.relevanceScore * 100).toFixed(0)}%
                                </Badge>
                              )}
                              {result.aiExplanation && (
                                <CheckCircle className="w-4 h-4 text-green-500" />
                              )}
                            </div>
                          </div>
                          <p className="text-xs text-gray-600 dark:text-gray-400 mb-2">
                            {result.snippet}
                          </p>
                          {result.keyReasons && (
                            <div className="text-xs">
                              <span className="font-medium">AI Analysis:</span>
                              <ul className="list-disc list-inside ml-2 mt-1">
                                {result.keyReasons.slice(0, 2).map((reason: string, i: number) => (
                                  <li key={i} className="text-gray-600 dark:text-gray-400">{reason}</li>
                                ))}
                              </ul>
                            </div>
                          )}
                        </Card>
                      ))}
                    </div>
                  </CardContent>
                </Card>
              )}
            </TabsContent>

            {/* Document Analysis Tab */}
            <TabsContent value="analysis" className="space-y-4">
              <div className="space-y-3">
                <Textarea
                  placeholder="Paste document content for AI analysis..."
                  value={analysisText}
                  onChange={(e) => setAnalysisText(e.target.value)}
                  className="min-h-32"
                />
                <div className="flex gap-2 flex-wrap">
                  {['summary', 'classification', 'key_points', 'quality_score'].map((type) => (
                    <Button 
                      key={type}
                      variant="outline" 
                      size="sm"
                      onClick={() => handleDocumentAnalysis(type)}
                      disabled={!analysisText.trim() || documentAnalysisMutation.isPending}
                    >
                      {documentAnalysisMutation.isPending ? (
                        <Loader2 className="w-3 h-3 animate-spin mr-1" />
                      ) : (
                        <FileText className="w-3 h-3 mr-1" />
                      )}
                      {type.replace('_', ' ').toUpperCase()}
                    </Button>
                  ))}
                </div>
              </div>

              {documentAnalysisMutation.data && (
                <Card className="bg-green-50 dark:bg-green-900/20 border-green-200 dark:border-green-800">
                  <CardHeader>
                    <CardTitle className="text-lg flex items-center gap-2">
                      <CheckCircle className="w-5 h-5 text-green-600" />
                      Analysis Result
                    </CardTitle>
                  </CardHeader>
                  <CardContent>
                    <div className="whitespace-pre-wrap text-sm">
                      {documentAnalysisMutation.data.analysis}
                    </div>
                  </CardContent>
                </Card>
              )}
            </TabsContent>

            {/* Research Synthesis Tab */}
            <TabsContent value="synthesis" className="space-y-4">
              <div className="space-y-3">
                <Input
                  placeholder="Research question for synthesis..."
                  value={query}
                  onChange={(e) => setQuery(e.target.value)}
                />
                <div className="flex items-center gap-2">
                  <span className="text-sm font-medium">Selected Documents:</span>
                  <Badge variant="outline">{selectedDocuments.length}</Badge>
                  <Button 
                    size="sm" 
                    variant="outline"
                    onClick={() => setSelectedDocuments([])}
                  >
                    Clear
                  </Button>
                </div>
                <Button 
                  onClick={handleResearchSynthesis}
                  disabled={!query.trim() || selectedDocuments.length === 0 || researchSynthesisMutation.isPending}
                  className="w-full"
                >
                  {researchSynthesisMutation.isPending ? (
                    <Loader2 className="w-4 h-4 animate-spin mr-2" />
                  ) : (
                    <Lightbulb className="w-4 h-4 mr-2" />
                  )}
                  Generate Research Synthesis
                </Button>
              </div>

              {researchSynthesisMutation.data && (
                <Card className="bg-purple-50 dark:bg-purple-900/20 border-purple-200 dark:border-purple-800">
                  <CardContent className="pt-4 space-y-4">
                    <div>
                      <h4 className="font-semibold text-purple-900 dark:text-purple-100 mb-2">Synthesis</h4>
                      <p className="text-sm">{researchSynthesisMutation.data.synthesis}</p>
                    </div>
                    
                    {researchSynthesisMutation.data.keyFindings.length > 0 && (
                      <div>
                        <h4 className="font-semibold text-purple-900 dark:text-purple-100 mb-2">Key Findings</h4>
                        <ul className="list-disc list-inside text-sm space-y-1">
                          {researchSynthesisMutation.data.keyFindings.map((finding: string, i: number) => (
                            <li key={i}>{finding}</li>
                          ))}
                        </ul>
                      </div>
                    )}

                    {researchSynthesisMutation.data.recommendations.length > 0 && (
                      <div>
                        <h4 className="font-semibold text-purple-900 dark:text-purple-100 mb-2">Recommendations</h4>
                        <ul className="list-disc list-inside text-sm space-y-1">
                          {researchSynthesisMutation.data.recommendations.map((rec: string, i: number) => (
                            <li key={i}>{rec}</li>
                          ))}
                        </ul>
                      </div>
                    )}
                  </CardContent>
                </Card>
              )}
            </TabsContent>

            {/* Embeddings Tab */}
            <TabsContent value="embeddings" className="space-y-4">
              <div className="space-y-3">
                <Input
                  placeholder="Text to generate embeddings..."
                  value={query}
                  onChange={(e) => setQuery(e.target.value)}
                />
                <Button 
                  onClick={handleGenerateEmbeddings}
                  disabled={!query.trim() || embeddingsMutation.isPending}
                  className="w-full"
                >
                  {embeddingsMutation.isPending ? (
                    <Loader2 className="w-4 h-4 animate-spin mr-2" />
                  ) : (
                    <Sparkles className="w-4 h-4 mr-2" />
                  )}
                  Generate Embeddings with Nebius
                </Button>
              </div>

              {embeddingsMutation.data && (
                <Card className="bg-yellow-50 dark:bg-yellow-900/20 border-yellow-200 dark:border-yellow-800">
                  <CardContent className="pt-4 space-y-3">
                    <div className="grid grid-cols-2 gap-4 text-sm">
                      <div>
                        <span className="font-medium">Model:</span>
                        <span className="ml-2">{embeddingsMutation.data.model}</span>
                      </div>
                      <div>
                        <span className="font-medium">Dimensions:</span>
                        <span className="ml-2">{embeddingsMutation.data.data[0].embedding.length}</span>
                      </div>
                    </div>
                    <div>
                      <span className="font-medium text-sm">Vector (first 10 dimensions):</span>
                      <div className="font-mono text-xs bg-white dark:bg-gray-800 p-2 rounded mt-1 overflow-x-auto">
                        [{embeddingsMutation.data.data[0].embedding.slice(0, 10).map((val: number) => val.toFixed(4)).join(', ')}...]
                      </div>
                    </div>
                    <div className="text-xs text-gray-600 dark:text-gray-400">
                      Token usage: {embeddingsMutation.data.usage.total_tokens} tokens
                    </div>
                  </CardContent>
                </Card>
              )}
            </TabsContent>
          </Tabs>
        </CardContent>
      </Card>

      {/* Error States */}
      {(aiSearchMutation.error || documentAnalysisMutation.error || researchSynthesisMutation.error || embeddingsMutation.error) && (
        <Card className="border-red-200 dark:border-red-800 bg-red-50 dark:bg-red-900/20">
          <CardContent className="pt-4">
            <div className="flex items-center gap-2 text-red-700 dark:text-red-300">
              <AlertCircle className="w-4 h-4" />
              <span className="font-medium">Error occurred</span>
            </div>
            <p className="text-sm text-red-600 dark:text-red-400 mt-1">
              {(aiSearchMutation.error || documentAnalysisMutation.error || researchSynthesisMutation.error || embeddingsMutation.error)?.message}
            </p>
          </CardContent>
        </Card>
      )}
    </div>
  );
}