import os
from pathlib import Path
from typing import Any

from pydantic import BaseModel
from pydantic_ai import RunContext

from haiku.rag.agents.research.models import Citation
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.rag.tools.qa import QAHistoryEntry
from haiku.skills.models import Skill, SkillSource
from haiku.skills.parser import parse_skill_md
from haiku.skills.state import SkillRunDeps

AGENT_PREAMBLE = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.

CRITICAL RULES:
1. For greetings or casual chat: respond directly WITHOUT using any tools
2. NEVER make up information - always use tools to get facts from the knowledge base
3. For questions: Use the "ask" tool - it handles search and citation automatically
4. For searches: Use the "search" tool - copy the ENTIRE tool response to your output INCLUDING content snippets
5. When you use the "ask" tool, summarize the key findings and always include citations in your response
"""


class ResearchEntry(BaseModel):
    question: str
    title: str
    executive_summary: str


class RAGState(BaseModel):
    citations: list[Citation] = []
    qa_history: list[QAHistoryEntry] = []
    document_filter: str | None = None
    searches: dict[str, list[SearchResult]] = {}
    documents: list[DocumentInfo] = []
    reports: list[ResearchEntry] = []


def create_skill(
    db_path: Path | None = None,
    config: Any = None,
) -> Skill:
    """Create a RAG skill for searching and analyzing documents.

    Args:
        db_path: Path to the LanceDB database. Resolved from:
            1. This argument
            2. HAIKU_RAG_DB environment variable
            3. haiku.rag default (config.storage.data_dir / "haiku.rag.lancedb")
        config: haiku.rag AppConfig instance. If None, uses get_config().
    """
    from haiku.rag.config import get_config

    if config is None:
        config = get_config()

    if db_path is None:
        env_db = os.environ.get("HAIKU_RAG_DB")
        if env_db:
            db_path = Path(env_db).expanduser()
        else:
            db_path = config.storage.data_dir / "haiku.rag.lancedb"

    path = Path(__file__).parent / "rag"
    metadata, instructions = parse_skill_md(path / "SKILL.md")

    async def _find_relevant_prior_qa(
        state: RAGState, query: str
    ) -> list[QAHistoryEntry]:
        from haiku.rag.embeddings import get_embedder
        from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD
        from haiku.rag.utils import cosine_similarity

        if not state.qa_history:
            return []

        embedder = get_embedder(config)
        query_embedding = await embedder.embed_query(query)

        to_embed = []
        to_embed_indices = []
        for i, qa in enumerate(state.qa_history):
            if qa.question_embedding is None:
                to_embed.append(qa.question)
                to_embed_indices.append(i)

        if to_embed:
            new_embeddings = await embedder.embed_documents(to_embed)
            for i, idx in enumerate(to_embed_indices):
                state.qa_history[idx].question_embedding = new_embeddings[i]

        matches = []
        for qa in state.qa_history:
            if qa.question_embedding is not None:
                similarity = cosine_similarity(query_embedding, qa.question_embedding)
                if similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD:
                    matches.append(qa)

        return matches

    async def search(
        ctx: RunContext[SkillRunDeps], query: str, limit: int | None = None
    ) -> str:
        """Search the knowledge base using hybrid search (vector + full-text).

        Returns ranked results with content and metadata.

        Args:
            query: The search query.
            limit: Maximum number of results.
        """
        from haiku.rag.client import HaikuRAG

        state = (
            ctx.deps.state
            if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState)
            else None
        )

        async with HaikuRAG(db_path, config=config, read_only=True) as rag:
            results = await rag.search(
                query,
                limit=limit,
                filter=state.document_filter if state else None,
            )
            results = await rag.expand_context(results)

        if state:
            state.searches[query] = list(results)

        return "\n\n---\n\n".join(
            r.format_for_agent(rank=i + 1, total=len(results))
            for i, r in enumerate(results)
        )

    async def list_documents(
        ctx: RunContext[SkillRunDeps],
        limit: int | None = None,
        offset: int | None = None,
        filter: str | None = None,
    ) -> list[dict[str, Any]]:
        """List documents in the knowledge base with optional pagination and filtering.

        Args:
            limit: Maximum number of documents to return.
            offset: Number of documents to skip.
            filter: Optional SQL WHERE clause to filter documents.
        """
        from haiku.rag.client import HaikuRAG

        async with HaikuRAG(db_path, config=config, read_only=True) as rag:
            documents = await rag.list_documents(limit, offset, filter)
            result = [
                {
                    "id": doc.id,
                    "title": doc.title,
                    "uri": doc.uri,
                    "metadata": doc.metadata,
                    "created_at": str(doc.created_at),
                    "updated_at": str(doc.updated_at),
                }
                for doc in documents
            ]

        if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
            for doc_dict in result:
                doc_info = DocumentInfo(
                    id=str(doc_dict["id"]),
                    title=doc_dict["title"] or "Untitled",
                    uri=doc_dict.get("uri") or "",
                    created=doc_dict.get("created_at", ""),
                )
                if not any(d.id == doc_info.id for d in ctx.deps.state.documents):
                    ctx.deps.state.documents.append(doc_info)

        return result

    async def get_document(
        ctx: RunContext[SkillRunDeps], query: str
    ) -> dict[str, Any] | None:
        """Retrieve a document by ID, title, or URI.

        Args:
            query: Document ID, title, or URI to look up.
        """
        from haiku.rag.client import HaikuRAG

        async with HaikuRAG(db_path, config=config, read_only=True) as rag:
            document = await rag.resolve_document(query)
            if document is None:
                return None
            result = {
                "id": document.id,
                "content": document.content,
                "title": document.title,
                "uri": document.uri,
                "metadata": document.metadata,
                "created_at": str(document.created_at),
                "updated_at": str(document.updated_at),
            }

        if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
            doc_info = DocumentInfo(
                id=str(result["id"]),
                title=result["title"] or "Untitled",
                uri=result.get("uri") or "",
                created=result.get("created_at", ""),
            )
            if not any(d.id == doc_info.id for d in ctx.deps.state.documents):
                ctx.deps.state.documents.append(doc_info)

        return result

    async def ask(ctx: RunContext[SkillRunDeps], question: str) -> str:
        """Ask a question and get an answer with citations from the knowledge base.

        Args:
            question: The question to ask.
        """
        from haiku.rag.client import HaikuRAG
        from haiku.rag.utils import format_citations

        state = (
            ctx.deps.state
            if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState)
            else None
        )

        ask_question = question
        if state:
            matches = await _find_relevant_prior_qa(state, question)
            if matches:
                prior_parts = []
                for qa in matches:
                    part = f"Q: {qa.question}\nA: {qa.answer}"
                    if qa.citations:
                        part += "\n" + format_citations(qa.citations)
                    prior_parts.append(part)
                ask_question = (
                    "Context from prior questions in this session:\n\n"
                    + "\n\n---\n\n".join(prior_parts)
                    + "\n\n---\n\nCurrent question: "
                    + question
                )

        async with HaikuRAG(db_path, config=config, read_only=True) as rag:
            answer, citations = await rag.ask(
                ask_question,
                filter=state.document_filter if state else None,
            )

        if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
            next_index = len(ctx.deps.state.citations) + 1
            for citation in citations:
                citation.index = next_index
                next_index += 1
            ctx.deps.state.citations.extend(citations)
            ctx.deps.state.qa_history.append(
                QAHistoryEntry(question=question, answer=answer, citations=citations)
            )

        if citations:
            answer += "\n\n" + format_citations(citations)

        return answer

    async def analyze(
        ctx: RunContext[SkillRunDeps],
        question: str,
        document: str | None = None,
        filter: str | None = None,
    ) -> str:
        """Answer complex analytical questions using code execution.

        Use this for questions requiring computation, aggregation, or
        data traversal across documents.

        Args:
            question: The question to answer.
            document: Optional document ID or title to pre-load for analysis.
            filter: Optional SQL WHERE clause to filter documents.
        """
        from haiku.rag.client import HaikuRAG

        async with HaikuRAG(db_path, config=config, read_only=True) as rag:
            documents = [document] if document else None
            result = await rag.rlm(question, documents=documents, filter=filter)
            output = result.answer
            if result.program:
                output += f"\n\nProgram:\n{result.program}"

        if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
            ctx.deps.state.qa_history.append(
                QAHistoryEntry(question=question, answer=output)
            )

        return output

    async def research(ctx: RunContext[SkillRunDeps], question: str) -> str:
        """Conduct deep multi-agent research on a question.

        Iteratively searches, analyzes, and synthesizes information from the
        knowledge base to produce a comprehensive research report.
        Only use when the user explicitly requests deep research.

        Args:
            question: The research question to investigate.
        """
        from haiku.rag.client import HaikuRAG

        state = (
            ctx.deps.state
            if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState)
            else None
        )

        async with HaikuRAG(db_path, config=config, read_only=True) as rag:
            report = await rag.research(
                question, filter=state.document_filter if state else None
            )

        if state:
            state.reports.append(
                ResearchEntry(
                    question=question,
                    title=report.title,
                    executive_summary=report.executive_summary,
                )
            )
            state.qa_history.append(
                QAHistoryEntry(question=question, answer=report.executive_summary)
            )

        parts = [
            f"# {report.title}",
            f"\n## Executive Summary\n{report.executive_summary}",
        ]
        if report.main_findings:
            parts.append("\n## Main Findings")
            for finding in report.main_findings:
                parts.append(f"- {finding}")
        if report.conclusions:
            parts.append("\n## Conclusions")
            for conclusion in report.conclusions:
                parts.append(f"- {conclusion}")
        if report.limitations:
            parts.append("\n## Limitations")
            for limitation in report.limitations:
                parts.append(f"- {limitation}")
        if report.recommendations:
            parts.append("\n## Recommendations")
            for rec in report.recommendations:
                parts.append(f"- {rec}")
        parts.append(f"\n## Sources\n{report.sources_summary}")

        return "\n".join(parts)

    return Skill(
        metadata=metadata,
        source=SkillSource.ENTRYPOINT,
        path=path,
        instructions=instructions,
        tools=[
            search,
            list_documents,
            get_document,
            ask,
            analyze,
            research,
        ],
        state_type=RAGState,
        state_namespace="rag",
    )