from collections.abc import Callable

from pydantic_ai import Agent

from pydstl.types import Evidence, SkillDocument

DISTILL_INSTRUCTIONS = """\
You are a knowledge distillation expert. Given a collection of evidence snippets, \
synthesize them into a clear, structured skill/guideline document.

IMPORTANT: Return the title separately in the `title` field. The `content` field must NOT \
include a title or heading — it should start directly with bullet points.

The content should:
- Be a flat list of bullet points, each prefixed with its topic like `topic: ...`
- Consolidate overlapping or redundant evidence
- Highlight key takeaways and actionable guidelines
- Cite or reference the source of evidence where relevant
- No headings, no sections, no nested structure — just a flat list of bullet points
"""

CONSOLIDATE_INSTRUCTIONS = """\
You are a knowledge consolidation expert. Given multiple skill/guideline documents \
(each distilled from a different source), merge them into a single cohesive document.

IMPORTANT: Return the title separately in the `title` field. The `content` field must NOT \
include a title or heading — it should start directly with bullet points.

The content should:
- Be a flat list of bullet points, each point prefixed with a concise topic title
- Deduplicate overlapping guidelines across sources
- Organize by theme, not by source
- Preserve specificity — don't water down concrete advice into generalities
- No headings, no sections, no nested structure — just a flat list of bullet points
"""

EDIT_INSTRUCTIONS = """\
You are a document editor. Given the current content of a skill document and an \
editing instruction, apply the requested changes and return the updated document.

Preserve the overall structure and tone unless the instruction says otherwise. \
Return the complete updated document content in markdown format.
"""


class LLM:
    def __init__(
        self,
        model: str,
        distill_fn: Callable | None = None,
        edit_fn: Callable | None = None,
        consolidate_fn: Callable | None = None,
    ):
        self._distill_fn = distill_fn
        self._edit_fn = edit_fn
        self._consolidate_fn = consolidate_fn
        self._model = model
        self._distill_agent: Agent[None, SkillDocument] | None = None
        self._edit_agent: Agent[None, str] | None = None
        self._consolidate_agent: Agent[None, SkillDocument] | None = None

    def _get_distill_agent(self) -> Agent[None, SkillDocument]:
        if self._distill_agent is None:
            self._distill_agent = Agent(
                model=self._model,
                output_type=SkillDocument,
                instructions=DISTILL_INSTRUCTIONS,
            )
        return self._distill_agent

    def _get_edit_agent(self) -> Agent[None, str]:
        if self._edit_agent is None:
            self._edit_agent = Agent(
                model=self._model,
                output_type=str,
                instructions=EDIT_INSTRUCTIONS,
            )
        return self._edit_agent

    def distill(self, evidence_list: list[Evidence], topic: str | None = None) -> SkillDocument:
        if self._distill_fn is not None:
            return self._distill_fn(evidence_list, topic)

        evidence_text = "\n\n---\n\n".join(
            f"**Evidence #{e.id}**"
            + (f" (source: {e.source})" if e.source else "")
            + f"\n{e.content}"
            for e in evidence_list
        )

        prompt = "Distill the following evidence into a skill document"
        if topic:
            prompt += f" focused on: {topic}"
        prompt += f"\n\n{evidence_text}"

        result = self._get_distill_agent().run_sync(prompt)
        return result.output

    def _get_consolidate_agent(self) -> Agent[None, SkillDocument]:
        if self._consolidate_agent is None:
            self._consolidate_agent = Agent(
                model=self._model,
                output_type=SkillDocument,
                instructions=CONSOLIDATE_INSTRUCTIONS,
            )
        return self._consolidate_agent

    def consolidate(
        self, documents: list[tuple[str, str]], topic: str | None = None
    ) -> SkillDocument:
        if self._consolidate_fn is not None:
            return self._consolidate_fn(documents, topic)

        docs_text = "\n\n===\n\n".join(
            f"**Document: {label}**\n\n{content}" for label, content in documents
        )

        prompt = "Consolidate the following skill documents into a single unified document"
        if topic:
            prompt += f" focused on: {topic}"
        prompt += f"\n\n{docs_text}"

        result = self._get_consolidate_agent().run_sync(prompt)
        return result.output

    def edit(self, current_content: str, instruction: str) -> str:
        if self._edit_fn is not None:
            return self._edit_fn(current_content, instruction)

        prompt = f"Current document:\n\n{current_content}\n\nEditing instruction: {instruction}"
        result = self._get_edit_agent().run_sync(prompt)
        return result.output
