"""Agentic research engine - coordinates agent skills and Ollama LLM interactions."""
from __future__ import annotations

import json
import os
import re
from typing import Any
import httpx

from app.agents.skills import SKILLS, Skill
from app.agents.tools import TOOLS
from app.config import get_settings


class Researcher:
    """Main research orchestrator that coordinates agents and tools."""

    def __init__(self):
        self.settings = get_settings()
        self.active_skills: list[Skill] = []
        self.findings: list[dict] = []
        self.context: dict = {}

    def select_skills(self, query: str, skill_names: list[str] | None = None) -> list[str]:
        """Select relevant agent skills for the query. Defaults to all if none specified."""
        if skill_names:
            selected = []
            for name in skill_names:
                s = SKILLS.get(name)
                if s:
                    selected.append(name)
            return selected

        query_lower = query.lower()
        selected = []
        for name, skill in SKILLS.items():
            for token in skill.verb_tokens:
                if token in query_lower:
                    selected.append(name)
                    break
        # If no token matched, default to researcher + qa
        if not selected:
            selected = ["researcher", "qa_agent"]
        return selected

    async def call_ollama(self, prompt: str, model: str | None = None, system: str | None = None) -> str:
        """Call Ollama LLM with streaming support."""
        model = model or self.settings.gpt_oss_model
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        async with httpx.AsyncClient(timeout=300) as client:
            resp = await client.post(
                f"{self.settings.ollama_url}/api/chat",
                json={"model": model, "messages": messages, "stream": False},
            )
            resp.raise_for_status()
            return resp.json().get("message", {}).get("content", "")

    async def run_skill(
        self, skill_name: str, query: str, doc_id: str | None = None
    ) -> str:
        """Execute a single agent skill."""
        skill = SKILLS.get(skill_name)
        if not skill:
            return f"[Unknown skill: {skill_name}]"

        system_prompt = f"""You are the {skill_name.replace('_', ' ').title()} agent.
{skill.instructions}"""

        tool_prompts = []
        for tool_name in skill.tools:
            tool_fn = TOOLS.get(tool_name)
            if not tool_fn:
                continue
            import inspect
            sig = inspect.signature(tool_fn)
            try:
                if tool_name == "read_document" and doc_id:
                    result = await tool_fn(doc_id)
                elif tool_name == "read_chunks" and doc_id:
                    result = await tool_fn(doc_id, limit=20)
                elif tool_name == "get_page_text" and doc_id:
                    result = await tool_fn(doc_id, page_num=0)
                elif tool_name == "extract_facts" and doc_id:
                    result = await tool_fn(doc_id, question=query)
                elif tool_name == "list_documents":
                    result = await tool_fn()
                elif tool_name in ("vector_search", "text_search", "memory_similarity_search"):
                    kwargs = {"query": query}
                    if doc_id and "doc_id" in sig.parameters:
                        kwargs["doc_id"] = doc_id
                    result = await tool_fn(**kwargs)
                elif tool_name in ("get_findings", "get_memories", "get_pipeline_state"):
                    result = await tool_fn(session_id=doc_id or "")
                elif tool_name == "save_finding":
                    result = await tool_fn(session_id=doc_id or "", question=query, answer="", finding_type="research", confidence=0.5)
                else:
                    kwargs = {}
                    if "doc_id" in sig.parameters and doc_id:
                        kwargs["doc_id"] = doc_id
                    if "query" in sig.parameters:
                        kwargs["query"] = query
                    result = await tool_fn(**kwargs)
                tool_prompts.append(f"\n--- {tool_name} output ---\n{result}")
            except Exception as e:
                tool_prompts.append(f"\n--- {tool_name} error ---\n{str(e)}")

        context = "".join(tool_prompts)
        
        prompt = f"""Research query: {query}
Document ID: {doc_id}
Context from tools:
{context}

Provide your analysis following the {skill_name} protocol."""

        response = await self.call_ollama(prompt, system=system_prompt)
        
        self.findings.append({
            "skill": skill_name,
            "query": query,
            "response": response,
            "doc_id": doc_id,
            "timestamp": str(os.popen('date +%Y-%m-%dT%H:%M:%S').read()).strip(),
        })
        
        return response

    async def run_research_session(
        self, query: str, session_id: str,
        doc_id: str | None = None,
        skill_names: list[str] | None = None
    ) -> dict[str, str]:
        """Run a full research session with selected agents."""
        selected_skills = self.select_skills(query, skill_names)
        results = {}
        
        for skill_name in selected_skills:
            results[skill_name] = await self.run_skill(skill_name, query, doc_id)
        
        return results

    async def cross_reference(self, query: str, doc_ids: list[str]) -> str:
        """Cross-reference content across multiple documents."""
        contexts = []
        for doc_id in doc_ids:
            chunks = await TOOLS["read_chunks"](doc_id)
            contexts.append(f"--- {doc_id} ---\n{chunks[:2000]}")
        
        combined = "\n\n".join(contexts)
        prompt = f"""Cross-reference these documents for the query: {query}

{combined}

Direct comparison. No preamble."""
        
        return await self.call_ollama(prompt, system="You are a cross-reference analyst. Output concise, comparative findings.")


# ── Pipeline skills ────────────────────────────────

PIPELINE_STAGES = ("document_triage", "evidence_extraction", "research_synthesis")


async def _parse_triage_output(response: str) -> dict:
    """Parse document-triage output into structured dict."""
    obj = {}
    for section in ["OBJECTIVE", "SUB-QUESTIONS", "CLASSIFICATION", "READING_ORDER", "EXTRACTION_CRITERIA", "RISKS_AND_GAPS"]:
        marker = f"[{section}]"
        # Find start of this section
        start = response.find(marker)
        end = response.find("[", start + len(marker)) if start != -1 else -1
        if end != -1:
            chunk = response[start + len(marker):end].strip()
        elif start != -1:
            chunk = response[start + len(marker):].strip()
        else:
            chunk = ""
        obj[section] = chunk
    return obj


async def _parse_evidence_output(response: str) -> list[dict]:
    """Parse evidence-extraction output into structured list."""
    rows = []
    # Find [EVIDENCE_ROWS] section
    start_marker = "[EVIDENCE_ROWS]"
    start = response.find(start_marker)
    # Check if there's a cross-comparison section
    cross_start = response.find("[CROSS_COMPARISON]")
    end = cross_start if cross_start != -1 else len(response)
    if start == -1:
        start = 0
    section = response[start:end].strip()
    
    lines = section.split("\n")
    for line in lines:
        line = line.strip()
        if not line or line.startswith("["):
            continue
        parts = [p.strip() for p in line.split("|")]
        if len(parts) >= 7:
            rows.append({
                "topic": parts[1],
                "evidence_type": parts[2],
                "description": parts[3],
                "trace_ref": parts[4],
                "evidence": parts[5],
                "analyst_note": parts[6],
                "confidence": parts[7] if len(parts) > 7 else "Medium",
                "review_needed": parts[8] if len(parts) > 8 else "No",
            })
    return rows


async def _parse_synthesis_mode(response: str) -> str:
    """Determine output mode from synthesis response content."""
    for mode in ["brief", "report", "gap analysis", "matrix"]:
        if mode in response.lower():
            return mode
    return "brief"


class ResearchPipeline:
    """Pipeline orchestrator: document_triage → evidence_extraction → research_synthesis."""

    def __init__(self):
        self.settings = get_settings()
        self.plan: dict = {}
        self.evidence: list[dict] = []
        self.synthesis: str = ""
        self.findings: list[dict] = []

    async def run(
        self,
        query: str,
        session_id: str,
        db,
        doc_ids: list[str] | None = None,
        output_mode: str | None = None,
    ) -> dict:
        """Run all three pipeline stages sequentially, persisting after each."""
        results = {}

        # ── Stage 1: document_triage ─────────────────────
        triage_skill = SKILLS["document_triage"]
        system = f"You are document_triage. {triage_skill.instructions}"
        
        # Gather context from available documents
        if doc_ids:
            doc_context = []
            for did in doc_ids:
                doc_info = await TOOLS["read_document"](did)
                doc_context.append(str(doc_info))
            context_input = "\n".join(doc_context)
        else:
            doc_list = await TOOLS["list_documents"]()
            context_input = doc_list
        
        triage_prompt = f"""Research query: {query}
Context from tools:
{context_input}

Apply the document-tireage protocol."""

        triage_output = await self._call_ollama(triage_prompt, system=system)
        self.plan = await _parse_triage_output(triage_output)
        self.findings.append({
            "stage": "triage",
            "output": triage_output,
            "state": self.plan,
        })
        await db.save_pipeline_stage(session_id, "triage", triage_output, self.plan)
        results["triage"] = triage_output
        results["triage_state"] = self.plan

        # ── Stage 2: evidence_extraction ────────────────
        evidence_skill = SKILLS["evidence_extraction"]
        system = f"""You are evidence_extraction. {evidence_skill.instructions}

Triage plan (from previous stage):
{json.dumps(self.plan, indent=2, default=str)}

Focus on extracting evidence for the sub-questions and extraction criteria defined above."""

        # Read chunks from all relevant documents
        all_chunks = []
        for did in (doc_ids or []):
            chunks = await TOOLS["read_chunks"](did)
            all_chunks.append(f"--- Document {did} ---\n{chunks}")
        
        context_input = "\n\n".join(all_chunks) if all_chunks else "No documents loaded yet."
        
        evidence_prompt = f"""Research query: {query}
Sub-questions to address:
{self.plan.get('SUB-QUESTIONS', 'N/A')}

Context from tools:
{context_input}

Apply the evidence_extraction protocol. Return structured evidence rows."""

        evidence_output = await self._call_ollama(evidence_prompt, system=system)
        self.evidence = await _parse_evidence_output(evidence_output)
        
        # Persist to DB
        if self.evidence:
            await db.save_structured_evidence(session_id, self.evidence)
        await db.save_pipeline_stage(session_id, "evidence", evidence_output, {"row_count": len(self.evidence)})
        
        results["evidence"] = evidence_output
        results["evidence_rows"] = self.evidence

        # ── Stage 3: research_synthesis ─────────────────
        synthesis_skill = SKILLS["research_synthesis"]
        
        mode = output_mode or await _parse_synthesis_mode(evidence_output)
        mode_prompts = {
            "brief": "Use Research Brief output mode.",
            "report": "Use Research Report output mode.",
            "gap": "Use Gap Analysis output mode.",
            "matrix": "Use Comparison Matrix output mode.",
        }
        mode_instruct = mode_prompts.get(mode, mode_prompts["brief"])
        
        system = f"""You are research_synthesis. {synthesis_skill.instructions}

{mode_instruct}

Extracted evidence (from previous stage):
{json.dumps(self.evidence, indent=2, default=str)[:15000]}

Apply the research_synthesis protocol."""

        synthesis_prompt = f"""Research query: {query}

Sub-questions:
{self.plan.get('SUB-QUESTIONS', 'N/A')}

Apply the research_synthesis protocol."""

        synthesis_output = await self._call_ollama(synthesis_prompt, system=system)
        
        self.synthesis = synthesis_output
        synthesis_state = {
            "mode": mode,
            "sub_questions": self.plan.get("SUB-QUESTIONS", ""),
            "evidence_count": len(self.evidence),
        }
        await db.save_pipeline_stage(session_id, "synthesis", synthesis_output, synthesis_state)
        
        results["synthesis"] = synthesis_output
        results["output_mode"] = mode
        results["pipeline_complete"] = True

        self.findings.append({
            "stage": "synthesis",
            "output": synthesis_output,
            "state": synthesis_state,
        })
        
        return results

    async def _call_ollama(self, prompt: str, system: str) -> str:
        """Call Ollama LLM."""
        messages = [
            {"role": "system", "content": system},
            {"role": "user", "content": prompt},
        ]
        async with httpx.AsyncClient(timeout=600) as client:
            resp = await client.post(
                f"{self.settings.ollama_url}/api/chat",
                json={"model": self.settings.gpt_oss_model, "messages": messages, "stream": False},
            )
            resp.raise_for_status()
            return resp.json().get("message", {}).get("content", "")


    async def render_pipeline_results(self, results: dict) -> list[dict]:
        """Render pipeline results for frontend display."""
        sections = []
        
        # Triage stage
        if "triage_state" in results:
            plan = results["triage_state"]
            sections.append({
                "stage": "triage",
                "title": "Stage 1: Source Triage",
                "objective": plan.get("OBJECTIVE", ""),
                "sub_questions": plan.get("SUB-QUESTIONS", ""),
                "classification": plan.get("CLASSIFICATION", ""),
                "reading_order": plan.get("READING_ORDER", ""),
                "extraction_criteria": plan.get("EXTRACTION_CRITERIA", ""),
                "risks_gaps": plan.get("RISKS_AND_GAPS", ""),
            })
        
        # Evidence stage
        if "evidence_rows" in results:
            rows = results["evidence_rows"]
            sections.append({
                "stage": "evidence",
                "title": f"Stage 2: Extracted Evidence ({len(rows)} rows)",
                "rows": rows,
            })
        
        # Synthesis stage
        if "synthesis" in results:
            sections.append({
                "stage": "synthesis",
                "title": f"Stage 3: Research Synthesis (mode: {results.get('output_mode', 'auto')})",
                "synthesis": results["synthesis"],
            })
        
        return sections