"""
AI Analyzer utility for document analysis using OpenAI.
Extracted from document_base.py AIAnalyzer class.
"""
import json
import re
from typing import Any, List, Dict, Optional

from app.services.llm_service import LLMService
from app.services.document_nodes import timer


def _parse_json_object_from_llm_text(text: str) -> Optional[Dict[str, Any]]:
    """
    Parse a JSON object from LLM text. Handles markdown fences (common from OIC / Anthropic)
    and balanced { ... } extraction when the model adds noise outside the object.
    """
    if text is None:
        return None
    t = str(text).strip()
    if not t:
        return None
    unfenced = re.sub(r"^\s*```(?:json)?\s*", "", t, flags=re.IGNORECASE | re.DOTALL)
    unfenced = re.sub(r"\s*```\s*$", "", unfenced, flags=re.DOTALL).strip()
    for body in (t, unfenced):
        if not body:
            continue
        try:
            data = json.loads(body)
            if isinstance(data, dict):
                return data
        except json.JSONDecodeError:
            pass
        start = body.find("{")
        if start < 0:
            continue
        depth = 0
        for i in range(start, len(body)):
            if body[i] == "{":
                depth += 1
            elif body[i] == "}":
                depth -= 1
                if depth == 0:
                    try:
                        data = json.loads(body[start : i + 1])
                        if isinstance(data, dict):
                            return data
                    except json.JSONDecodeError:
                        pass
                    break
    return None


_CHUNK_CITATION_PAREN_RE = re.compile(
    r"\s*\((?:Chunks?|Source chunks?|chunk ids?)[^)]*\)",
    re.IGNORECASE,
)


def _clean_user_facing_document_text(text: str) -> str:
    """Remove chunk-id footnotes the model sometimes adds to @file summaries."""
    if not text:
        return text
    cleaned = _CHUNK_CITATION_PAREN_RE.sub("", str(text))
    cleaned = re.sub(r"\s*,\s*chunk_\w+", "", cleaned, flags=re.IGNORECASE)
    cleaned = re.sub(r"\bchunk_\w+\b,?\s*", "", cleaned, flags=re.IGNORECASE)
    return re.sub(r"\n{3,}", "\n\n", cleaned).strip()


def _unwrap_find_in_document_nested_answer(d: Dict[str, Any]) -> Dict[str, Any]:
    """
    If the model put a second JSON object (optionally in ```json fences) inside ``answer``,
    merge the inner ``found`` / ``answer`` / ``quotes`` / ``confidence`` into the top level.
    """
    out: Dict[str, Any] = dict(d)
    for _ in range(3):
        ans = out.get("answer")
        if not isinstance(ans, str):
            break
        a = ans.strip()
        if not a:
            break
        inner = _parse_json_object_from_llm_text(a)
        if not isinstance(inner, dict) or "answer" not in inner:
            break
        for k in ("found", "answer", "quotes", "confidence"):
            if k in inner:
                out[k] = inner[k]
        out.pop("raw_response", None)
    out.pop("raw_response", None)
    return out


class AIAnalyzer:
    """
    Utility class for AI-powered document analysis using OpenAI.
    """
    
    def __init__(self, api_key: str, claude_api_key: str = ""):
        self.api_key = api_key
        self._llm = (
            LLMService(api_key=api_key or "", claude_api_key=claude_api_key or "")
            if (api_key or "").strip() or (claude_api_key or "").strip()
            else None
        )
    
    def extract_folder_from_query(self, query: str, 
                                   available_folders: List[str]) -> Optional[str]:
        """Extract folder reference from natural language query"""
        if not self._llm or not available_folders:
            return None
        
        try:
            prompt = f"""Analyze this query and determine if it mentions a specific folder.
Available folders: {', '.join(available_folders)}

Query: "{query}"

If the query mentions a folder (like "in documents", "from reports", "in the X folder"), 
return the matching folder name. Otherwise return "none".

Return ONLY the folder name or "none", nothing else."""

            with timer("ai.extract_folder_from_query"):
                result = self._llm.chat_completion(
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=50,
                    temperature=0,
                    generation_name="extract_folder_from_query",
                ).strip().lower()
            
            if result != "none":
                for folder in available_folders:
                    if folder.lower() in result or result in folder.lower():
                        return folder
            
            return None
            
        except Exception as e:
            print(f"Error extracting folder: {e}")
            return None
    
    def analyze_document(self, content: str, filename: str, langfuse_trace=None) -> Dict:
        """Analyze document content and extract key information"""
        if not self._llm:
            return {"error": "OpenAI client not initialized"}
        
        try:
            # Truncate content if too long
            max_content = content[:8000] if len(content) > 8000 else content
            
            prompt = f"""You are a document analysis assistant. Analyze the following document and extract structured information.

Filename: {filename}
Content:
{max_content}

You MUST respond with ONLY a valid JSON object with NO additional text or explanation. The JSON must have this exact structure:

{{
  "summary": "A brief 2-3 sentence overview of the document",
  "key_points": ["Point 1", "Point 2", "Point 3", "Point 4", "Point 5"],
  "topics": ["Topic 1", "Topic 2", "Topic 3"],
  "entities": {{
    "names": ["Person 1", "Organization 1"],
    "dates": ["Date 1", "Date 2"],
    "numbers": ["Amount 1", "Quantity 1"]
  }}
}}

Rules:
1. Respond ONLY with the JSON object, no markdown code blocks, no explanations
2. key_points must be an array of strings (max 5 items)
3. topics must be an array of strings
4. entities must be an object with names, dates, and numbers arrays
5. If any field has no data, use an empty array []
6. Ensure all strings are properly escaped for valid JSON
7. Never wrap the JSON in markdown (no ``` or ```json)."""

            with timer("ai.analyze_document"):
                result = self._llm.chat_completion(
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=1000,
                    temperature=0.3,
                    langfuse_trace=langfuse_trace,
                    generation_name="analyze_document",
                ).strip()
            
            parsed = _parse_json_object_from_llm_text(result)
            if isinstance(parsed, dict):
                return parsed
            return {"summary": result, "raw_response": True}
                
        except Exception as e:
            return {"error": str(e)}
    
    def summarize_search_results(self, query: str, results: List[Dict],
                                  folder_detected: str = None, langfuse_trace=None) -> str:
        """Summarize search results into a coherent answer"""
        if not self._llm:
            return "Error: OpenAI client not initialized"
        
        if not results:
            return "No relevant documents found for your query."
        
        try:
            # Build context from chunk-level hits (ranked).
            context_parts = []
            for i, result in enumerate(results[:5]):
                chunk_id = result.get("chunk_id") or result.get("id", "")
                filename = result.get("filename", "Unknown")
                rel = result.get("relevance")
                rel_note = f" (relevance {rel:.1f}%)" if isinstance(rel, (int, float)) else ""
                pages = result.get("page_numbers") or result.get("page_span") or ""
                page_note = f" pages {pages}" if pages else ""
                doc_info = f"Chunk {chunk_id} from {filename}{rel_note}{page_note}"
                content = (result.get("chunk_text") or result.get("content") or "")[:2000]
                context_parts.append(f"{doc_info}\n{content}")
            
            context = "\n\n---\n\n".join(context_parts)
            
            folder_note = f" (searched in folder: {folder_detected})" if folder_detected else ""
            
            prompt = f"""Based on the following document excerpts, answer this question{folder_note}:

Question: {query}

Excerpts (chunk-level):
{context}

Provide a comprehensive, factual answer based on the excerpts for a business user.

IMPORTANT RULES:
- Do NOT cite chunk IDs, chunk names (e.g. chunk_12), or parenthetical chunk lists like "(Chunks: ...)".
- You may mention a source document filename once when helpful; otherwise write standalone prose.
- Do NOT include any conversational filler, offers, or follow-up suggestions.
- Do NOT say things like "If you need the actual PDF...", "Let me know if...", "Feel free to ask...", "Please specify...", or similar phrases.
- Do NOT offer to provide more details or breakdowns.
- Your response must be a complete, self-contained answer with no interactive elements.
- End your response with the final piece of information — do NOT add a closing offer or question."""

            with timer("ai.summarize_search_results"):
                raw = self._llm.chat_completion(
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=1500,
                    temperature=0.3,
                    langfuse_trace=langfuse_trace,
                    generation_name="summarize_search_results",
                ).strip()
                return _clean_user_facing_document_text(raw)
            
        except Exception as e:
            return f"Error generating summary: {str(e)}"
    
    def generate_key_findings(self, query: str, results: List[Dict],
                               folder_detected: str = None, langfuse_trace=None) -> Dict:
        """Generate key findings from search results - similar to analyze_document but for multiple docs"""
        if not self._llm:
            return {"error": "OpenAI client not initialized"}
        
        if not results:
            return {
                "summary": "No relevant documents found for your query.",
                "key_points": []
            }
        
        try:
            context_parts = []
            for i, result in enumerate(results[:5]):
                chunk_id = result.get("chunk_id") or result.get("id", "")
                filename = result.get("filename", "Unknown")
                rel = result.get("relevance")
                rel_note = f" (relevance {rel:.1f}%)" if isinstance(rel, (int, float)) else ""
                pages = result.get("page_numbers") or result.get("page_span") or ""
                page_note = f" pages {pages}" if pages else ""
                doc_info = f"Chunk {chunk_id} from {filename}{rel_note}{page_note}"
                content = (result.get("chunk_text") or result.get("content") or "")[:2000]
                context_parts.append(f"{doc_info}\n{content}")
            
            context = "\n\n---\n\n".join(context_parts)
            folder_note = f" (from folder: {folder_detected})" if folder_detected else ""
            
            prompt = f"""You are a document analysis assistant. Analyze the following search results for the query and extract key findings.

Query: {query}{folder_note}

Excerpts (chunk-level):
{context}

You MUST respond with ONLY a valid JSON object with NO additional text or explanation. The JSON must have this exact structure:

{{
  "summary": "A brief 2-3 sentence summary answering the query based on documents",
  "key_points": ["Key finding 1", "Key finding 2", "Key finding 3", "Key finding 4", "Key finding 5"]
}}


Rules:
1. Respond ONLY with the JSON object, no markdown code blocks, no explanations
2. key_points must be an array of strings (max 5 most important findings)
3. Do NOT mention chunk IDs or "(Chunks: ...)" in summary or key_points
4. If any field has no data, use an empty array []
5. Ensure all strings are properly escaped for valid JSON
6. Never wrap the JSON in markdown (no ``` or ```json)."""

            with timer("ai.generate_key_findings"):
                result = self._llm.chat_completion(
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=1000,
                    temperature=0.3,
                    langfuse_trace=langfuse_trace,
                    generation_name="generate_key_findings",
                ).strip()
            
            parsed = _parse_json_object_from_llm_text(result)
            if isinstance(parsed, dict):
                if isinstance(parsed.get("summary"), str):
                    parsed["summary"] = _clean_user_facing_document_text(parsed["summary"])
                if isinstance(parsed.get("key_points"), list):
                    parsed["key_points"] = [
                        _clean_user_facing_document_text(p)
                        for p in parsed["key_points"]
                        if isinstance(p, str)
                    ]
                return parsed
            return {
                "summary": _clean_user_facing_document_text(result),
                "key_points": [],
                "raw_response": True,
            }
                
        except Exception as e:
            return {"error": str(e), "key_points": []}
    
    def find_in_document(self, query: str, content: str,
                         filename: str, langfuse_trace=None) -> Dict:
        """Find specific information within a document"""
        if not self._llm:
            return {"error": "OpenAI client not initialized"}
        
        try:
            max_content = content[:12000] if len(content) > 12000 else content
            
            prompt = f"""Search this document for information related to the query.

Document: {filename}
Query: {query}

Content:
{max_content}

Return a single JSON object (and nothing else) with:
- found: boolean indicating if relevant info was found
- answer: plain text — the direct answer, not JSON inside a string
- quotes: list of relevant short quotes from the document (max 3 strings)
- confidence: one of "high", "medium", or "low"

Rules:
- Output raw JSON only. Do not use markdown code fences (no ``` or ```json).
- Do not put a second JSON object inside the "answer" field — the answer must be human-readable text."""

            with timer("ai.find_in_document"):
                result = self._llm.chat_completion(
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=1000,
                    temperature=0.2,
                    langfuse_trace=langfuse_trace,
                    generation_name="find_in_document",
                ).strip()
            
            parsed = _parse_json_object_from_llm_text(result)
            if isinstance(parsed, dict):
                return _unwrap_find_in_document_nested_answer(parsed)
            return {"found": True, "answer": result, "raw_response": True}
                
        except Exception as e:
            return {"error": str(e)}

    def decompose_query(self, query: str, langfuse_trace=None) -> List[str]:
        """
        Split a complex question into 2–3 focused sub-queries.
        Returns [query] on any failure.
        """
        if not self._llm:
            return [query]
        q = (query or "").strip()
        if not q:
            return [query]

        prompt = f"""Break the following question into 2-3 independent sub-questions that can be answered separately.
Return ONLY valid JSON of the form:
{{"sub_queries": ["...","..."]}}

Question: {q}
"""
        try:
            with timer("ai.decompose_query"):
                raw = self._llm.chat_completion(
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=300,
                    temperature=0.2,
                    langfuse_trace=langfuse_trace,
                    generation_name="decompose_query",
                ).strip()
            parsed = json.loads(raw)
            subs = parsed.get("sub_queries", [])
            subs = [s.strip() for s in subs if isinstance(s, str) and s.strip()]
            return subs[:3] if subs else [q]
        except Exception:
            return [q]
