"""Core types for the prompt builder framework.

Every component of a prompt is a typed, validated Pydantic model.
The Prompt object is the source tree that the compiler consumes.
"""

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any

from pydantic import BaseModel, Field


# ── Enums ──────────────────────────────────────────────────────────────────


class Scope(str, Enum):
    ALWAYS = "always"


class RulePriority(str, Enum):
    HARD = "hard"  # never cut
    SOFT = "soft"  # can be dropped under pressure


class CoTStrategy(str, Enum):
    NONE = "none"
    STEP_BY_STEP = "step_by_step"
    STRUCTURED = "structured"
    REACT = "react"


class Visibility(str, Enum):
    HIDDEN = "hidden"
    VISIBLE = "visible"
    EXTRACTED = "extracted"


class Section(str, Enum):
    PERSONA = "persona"
    RULES = "rules"
    INSTRUCTION = "instruction"
    OUTPUT_SPEC = "output_spec"
    COT = "cot"
    EXAMPLES = "examples"
    CONTEXT = "context"
    HISTORY = "history"
    INPUT_DATA = "input_data"


class TruncationStrategy(str, Enum):
    NEVER_CUT = "never_cut"
    CUT_BY_RELEVANCE = "cut_by_relevance"
    CUT_BY_WEIGHT = "cut_by_weight"
    TRUNCATE_OLDEST = "truncate_oldest"
    CUT_LAST = "cut_last"


# ── Identity Layer ─────────────────────────────────────────────────────────


class Persona(BaseModel):
    """Who the LLM is."""

    role: str
    tone: str | None = None
    domain: list[str] = Field(default_factory=list)
    anti_patterns: list[str] = Field(default_factory=list)


class WhenScope(BaseModel):
    """Conditional scope for rules."""

    condition: str


class Rule(BaseModel):
    """A hard constraint / invariant."""

    constraint: str
    scope: Scope | WhenScope = Scope.ALWAYS
    priority: RulePriority = RulePriority.HARD


# ── Knowledge Layer ────────────────────────────────────────────────────────


class ContextBlock(BaseModel):
    """A chunk of injected knowledge (RAG, docs, memory)."""

    content: str
    source: str = "manual"
    relevance: float = Field(default=1.0, ge=0.0, le=1.0)
    token_count: int | None = None  # auto-computed by compiler if None
    metadata: dict[str, Any] = Field(default_factory=dict)
    freshness: datetime | None = None


class Example(BaseModel):
    """A few-shot demonstration."""

    input: str
    output: str
    explanation: str | None = None
    negative: bool = False
    weight: float = Field(default=1.0, ge=0.0)


# ── Task Layer ─────────────────────────────────────────────────────────────


class Instruction(BaseModel):
    """What to do and why."""

    task: str
    subtasks: list[str] = Field(default_factory=list)
    constraints: list[str] = Field(default_factory=list)
    intent: str | None = None


class InputData(BaseModel):
    """The thing to operate on."""

    content: str | dict | list
    format: str = "plain"
    delimiter: str = "```"
    label: str = "input"


# ── Control Layer ──────────────────────────────────────────────────────────


class OutputSpec(BaseModel):
    """How the output should look."""

    format: str = "plain"
    schema_model: type | None = Field(default=None, exclude=True)
    length: str | None = None
    language: str | None = None
    constraints: list[str] = Field(default_factory=list)

    model_config = {"arbitrary_types_allowed": True}


class CoTSpec(BaseModel):
    """Chain-of-thought configuration."""

    strategy: CoTStrategy = CoTStrategy.NONE
    visibility: Visibility = Visibility.VISIBLE
    format: str | None = None


# ── Conversation Layer ─────────────────────────────────────────────────────


class Message(BaseModel):
    """A conversation turn."""

    role: str  # "user", "assistant", "tool"
    content: str
    token_count: int | None = None  # auto-computed by compiler if None
    timestamp: datetime | None = None


# ── Meta Layer ─────────────────────────────────────────────────────────────


class PriorityEntry(BaseModel):
    """A single entry in the priority map."""

    section: Section
    strategy: TruncationStrategy


# Default priority order: what gets cut first → last
DEFAULT_PRIORITY: list[PriorityEntry] = [
    PriorityEntry(section=Section.CONTEXT, strategy=TruncationStrategy.CUT_BY_RELEVANCE),
    PriorityEntry(section=Section.HISTORY, strategy=TruncationStrategy.TRUNCATE_OLDEST),
    PriorityEntry(section=Section.EXAMPLES, strategy=TruncationStrategy.CUT_BY_WEIGHT),
    PriorityEntry(section=Section.COT, strategy=TruncationStrategy.CUT_LAST),
    PriorityEntry(section=Section.PERSONA, strategy=TruncationStrategy.NEVER_CUT),
    PriorityEntry(section=Section.RULES, strategy=TruncationStrategy.NEVER_CUT),
    PriorityEntry(section=Section.INSTRUCTION, strategy=TruncationStrategy.NEVER_CUT),
    PriorityEntry(section=Section.OUTPUT_SPEC, strategy=TruncationStrategy.NEVER_CUT),
    PriorityEntry(section=Section.INPUT_DATA, strategy=TruncationStrategy.NEVER_CUT),
]


class PriorityMap(BaseModel):
    """Declares what gets cut and in what order when tokens are tight."""

    entries: list[PriorityEntry] = Field(default_factory=lambda: list(DEFAULT_PRIORITY))

    def get_strategy(self, section: Section) -> TruncationStrategy:
        for entry in self.entries:
            if entry.section == section:
                return entry.strategy
        return TruncationStrategy.NEVER_CUT

    def cutting_order(self) -> list[PriorityEntry]:
        """Sections in order of what gets cut first."""
        return [e for e in self.entries if e.strategy != TruncationStrategy.NEVER_CUT]


# ── Token Budget ───────────────────────────────────────────────────────────


class TokenBudget(BaseModel):
    """How many tokens we have to work with."""

    context_window: int
    output_reserve: int = 4096
    safety_margin: int = 200

    @property
    def available(self) -> int:
        return self.context_window - self.output_reserve - self.safety_margin


# ── The Prompt ─────────────────────────────────────────────────────────────


class Prompt(BaseModel):
    """The complete prompt source tree.

    This is the input to the compiler. It contains every piece of information
    that could go into the final message array. The compiler decides what fits.
    """

    # Identity
    persona: Persona | None = None
    rules: list[Rule] = Field(default_factory=list)

    # Knowledge
    context: list[ContextBlock] = Field(default_factory=list)
    examples: list[Example] = Field(default_factory=list)

    # Task
    instruction: Instruction
    input_data: InputData | None = None

    # Control
    output_spec: OutputSpec | None = None
    cot: CoTSpec | None = None

    # Conversation
    history: list[Message] = Field(default_factory=list)

    # Meta
    priority: PriorityMap = Field(default_factory=PriorityMap)
    version: str = "0.1.0"

    model_config = {"arbitrary_types_allowed": True}


# ── Compilation Result ─────────────────────────────────────────────────────


class SectionReport(BaseModel):
    """What happened to a single section during compilation."""

    section: Section
    included_tokens: int = 0
    original_tokens: int = 0
    items_included: int = 0
    items_dropped: int = 0
    strategy_applied: TruncationStrategy | None = None


class CompilationResult(BaseModel):
    """The output of the compiler."""

    messages: list[dict[str, str]]
    sections: list[SectionReport] = Field(default_factory=list)
    total_tokens: int = 0
    budget_available: int = 0
    warnings: list[str] = Field(default_factory=list)

    model_config = {"arbitrary_types_allowed": True}

    @property
    def headroom(self) -> int:
        return self.budget_available - self.total_tokens

    def report(self) -> str:
        """Human-readable compilation report."""
        lines = [
            f"Compilation Report (v{self.total_tokens}/{self.budget_available} tokens)",
            f"  Headroom: {self.headroom} tokens",
            "",
        ]
        for s in self.sections:
            status = "✓" if s.items_dropped == 0 else "⚠"
            lines.append(
                f"  {status} {s.section.value:12s}  "
                f"{s.included_tokens:>6d} tokens  "
                f"({s.items_included} kept, {s.items_dropped} dropped)"
            )
        if self.warnings:
            lines.append("")
            for w in self.warnings:
                lines.append(f"  ⚠ {w}")
        return "\n".join(lines)
