import os
from functools import cache
from pathlib import Path
from typing import Any

from pydantic import BaseModel, Field
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, SkillMetadata, SkillSource, StateMetadata
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 skills to get facts from the knowledge base
3. When a skill returns citations, always include them in your response
"""


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


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


STATE_TYPE = RAGState
STATE_NAMESPACE = "rag"

_skill_path = Path(__file__).parent / "rag"


@cache
def skill_metadata() -> SkillMetadata:
    metadata, _ = parse_skill_md(_skill_path / "SKILL.md")
    return metadata


@cache
def instructions() -> str | None:
    _, instr = parse_skill_md(_skill_path / "SKILL.md")
    return instr


def state_metadata() -> StateMetadata:
    return StateMetadata(
        namespace=STATE_NAMESPACE,
        type=STATE_TYPE,
        schema=STATE_TYPE.model_json_schema(),
    )


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"

    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,
    ) -> list[dict[str, Any]]:
        """List documents in the knowledge base with optional pagination.

        Args:
            limit: Maximum number of documents to return.
            offset: Number of documents to skip.
        """
        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)
            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 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=skill_metadata(),
        source=SkillSource.ENTRYPOINT,
        path=_skill_path,
        instructions=instructions(),
        tools=[
            search,
            list_documents,
            get_document,
            ask,
            research,
        ],
        state_type=STATE_TYPE,
        state_namespace=STATE_NAMESPACE,
    )