"""Compiler core functions for Lattice.

Assembles prompts, parses LLM output, extracts proposals, and manages session truncation.
All functions are pure (no I/O) — Shell layer handles file operations.

Reference: RFC-002 §4.3 Project Evolution
"""

from __future__ import annotations

import hashlib
import re
from typing import TYPE_CHECKING

import deal

from lattice.core.types.evidence import CotPhases
from lattice.core.types.log import LogEntry, Session
from lattice.core.types.proposal import PatternType, Proposal, ProposalAction
from lattice.core.types.rule import Rule, RuleDiff

if TYPE_CHECKING:
    pass


@deal.pre(lambda rules: isinstance(rules, list))
@deal.pre(lambda rules: all(isinstance(r, Rule) for r in rules))
@deal.post(lambda result: isinstance(result, str))
def _format_rules(rules: list[Rule]) -> str:
    """Format rules for prompt context."""
    if not rules:
        return "No existing rules."
    lines = []
    for rule in rules:
        lines.append(f"### {rule.title}")
        lines.append(f"File: {rule.file_path}")
        lines.append(rule.content)
        lines.append("")
    return "\n".join(lines)


@deal.pre(lambda sessions: isinstance(sessions, list))
@deal.pre(lambda sessions: all(isinstance(s, Session) for s in sessions))
@deal.post(lambda result: isinstance(result, str))
def _format_sessions(sessions: list[Session]) -> str:
    """Format sessions for prompt context."""
    if not sessions:
        return "No sessions to analyze."
    lines = []
    for session in sessions:
        lines.append(f"## Session: {session.session_id}")
        for log in session.logs:
            lines.append(f"[{log.role.value}]: {log.content}")
        lines.append("")
    return "\n".join(lines)


@deal.pre(
    lambda current_token_count, alert_tokens: (
        isinstance(current_token_count, int) and isinstance(alert_tokens, int)
    )
)
@deal.post(lambda result: isinstance(result, str))
def _format_token_warning(current_token_count: int, alert_tokens: int) -> str:
    """Format token warning message."""
    if current_token_count > alert_tokens:
        return f"""
**WARNING**: Current token count ({current_token_count}) exceeds alert threshold ({alert_tokens}).
Be aggressive in consolidation. Merge overlapping rules. Remove rules that are no longer relevant.
Prefis fewer, stronger rules over many weak ones.
"""
    return ""


@deal.pre(
    lambda sessions, current_rules, alert_tokens, current_token_count, prompt_template: (
        isinstance(sessions, list)
        and isinstance(current_rules, list)
        and isinstance(alert_tokens, int)
        and isinstance(current_token_count, int)
        and isinstance(prompt_template, str)
        and all(isinstance(s, Session) for s in sessions)
        and all(isinstance(r, Rule) for r in current_rules)
        and alert_tokens > 0
        and current_token_count >= 0
        and len(prompt_template) > 0
        and "<triage>" in prompt_template
        and "<cross_ref>" in prompt_template
        and "<synthesis>" in prompt_template
        and "<review>" in prompt_template
    )
)
@deal.post(lambda result: isinstance(result, str))
@deal.post(lambda result: "<triage>" in result)
@deal.post(lambda result: "<cross_ref>" in result)
@deal.post(lambda result: "<synthesis>" in result)
@deal.post(lambda result: "<review>" in result)
def assemble_compiler_prompt(
    sessions: list[Session],
    current_rules: list[Rule],
    alert_tokens: int,
    current_token_count: int,
    prompt_template: str,
) -> str:
    """Assemble the full prompt for the Project Compiler LLM call.

    Loads the prompt template and interpolates:
    - Current rules as formatted context
    - Session logs (pre-truncated)
    - Four-phase CoT instructions
    - Token warning if current_token_count > alert_tokens

    Args:
        sessions: List of Session objects to analyze.
        current_rules: List of current Rule objects.
        alert_tokens: Token threshold for aggressive consolidation warning.
        current_token_count: Current total token count of rules.
        prompt_template: Template string with {current_rules}, {sessions},
            {alert_tokens}, {current_token_count}, {token_warning} placeholders.

    Returns:
        Complete prompt string containing all four phase tags.

    >>> # Template must contain all four phase tags for @post contracts
    >>> result = assemble_compiler_prompt([], [], 3000, 0, "Rules: {current_rules}<triage><cross_ref><synthesis><review>")
    >>> "Rules: No existing rules." in result
    True
    >>> result = assemble_compiler_prompt([], [], 3000, 0, "<triage><cross_ref><synthesis><review>")
    >>> all(tag in result for tag in ["<triage>", "<cross_ref>", "<synthesis>", "<review>"])
    True
    """
    formatted_rules = _format_rules(current_rules)
    formatted_sessions = _format_sessions(sessions)
    token_warning = _format_token_warning(current_token_count, alert_tokens)

    return prompt_template.format(
        current_rules=formatted_rules,
        sessions=formatted_sessions,
        alert_tokens=str(alert_tokens),
        current_token_count=str(current_token_count),
        token_warning=token_warning,
    )


@deal.pre(lambda raw_output: isinstance(raw_output, str))
@deal.post(lambda result: isinstance(result, CotPhases))
def parse_cot_output(raw_output: str) -> CotPhases:
    """Parse LLM output into structured CoT phases.

    Extracts text between XML-style tags:
    <triage>...</triage>, <cross_ref>...</cross_ref>,
    <synthesis>...</synthesis>, <review>...</review>

    Tolerates missing phases (returns empty string for missing phase).
    Tolerates extra whitespace and minor formatting variations.

    Args:
        raw_output: Raw LLM output string.

    Returns:
        CotPhases with extracted phase content (empty string for missing phases).

    >>> phases = parse_cot_output("<triage>content</triage>")
    >>> phases.triage
    'content'
    >>> phases = parse_cot_output("<triage>a</triage><review>b</review>")
    >>> phases.cross_ref
    ''
    >>> phases.review
    'b'
    """
    tags = ["triage", "cross_ref", "synthesis", "review"]
    extracted = {}

    for tag in tags:
        pattern = rf"<{tag}>(.*?)</{tag}>"
        match = re.search(pattern, raw_output, re.DOTALL)
        extracted[tag] = match.group(1).strip() if match else ""

    return CotPhases(
        triage=extracted["triage"],
        cross_ref=extracted["cross_ref"],
        synthesis=extracted["synthesis"],
        review=extracted["review"],
    )


# @invar:allow function_size: Proposal parsing logic is cohesive; extraction would reduce clarity
@deal.pre(lambda phases: isinstance(phases, CotPhases))
@deal.post(lambda result: isinstance(result, list))
@deal.post(lambda result: all(isinstance(p, Proposal) for p in result))
def extract_proposals(phases: CotPhases) -> list[Proposal]:
    """Extract structured proposals from synthesis and review phases.

    Parses the <synthesis> and <review> phase content for:
    - New rule additions
    - Rule merges
    - Rule removals
    - Rule rewordings

    Each proposal must have: action type, target file, content, evidence, rationale.

    Args:
        phases: CotPhases object with synthesis and review content.

    Returns:
        List of Proposal objects. May be empty (zero-proposal run).

    >>> from lattice.core.types.evidence import CotPhases
    >>> phases = CotPhases(triage="", cross_ref="", synthesis="", review="")
    >>> extract_proposals(phases)
    []
    """
    proposals: list[Proposal] = []

    # Combine synthesis and review for proposal extraction
    combined = f"{phases.synthesis}\n{phases.review}"

    # Pattern for PROPOSAL blocks
    proposal_pattern = re.compile(
        r"##\s*PROPOSAL:\s*(ADD|MERGE|REMOVE|REWORD)\s*-\s*(.+?)\n"
        r"Action:\s*(ADD|MERGE|REMOVE|REWORD)\s*\n"
        r"Target:\s*(.+?)\s*\n"
        r"Content:\s*\|?\s*\n(.*?)(?=Evidence:|$)"
        r"(?:Evidence:\s*(.+?)\s*\n)?"
        r"(?:Rationale:\s*\|?\s*\n(.*?)(?=##|$))?",
        re.DOTALL | re.IGNORECASE,
    )

    for match in proposal_pattern.finditer(combined):
        action_str = match.group(3).strip().upper()
        title = match.group(2).strip()
        target = match.group(4).strip()
        content = match.group(5).strip()
        evidence_str = match.group(6) or ""
        rationale = (match.group(7) or "").strip()

        # Map action string to enum
        action_map = {
            "ADD": ProposalAction.ADD,
            "MERGE": ProposalAction.MERGE,
            "REMOVE": ProposalAction.REMOVE,
            "REWORD": ProposalAction.UPDATE,  # REWORD maps to UPDATE
            "UPDATE": ProposalAction.UPDATE,
        }
        action = action_map.get(action_str, ProposalAction.ADD)

        # Parse evidence session IDs
        evidence_ids = []
        for evidence_match in re.finditer(
            r"session[-_]?(\w+)", evidence_str, re.IGNORECASE
        ):
            evidence_ids.append(evidence_match.group(0))

        # Generate proposal ID from title (hashlib imported at module level)
        proposal_id = hashlib.md5(title.encode()).hexdigest()[:12]

        proposal = Proposal(
            proposal_id=proposal_id,
            pattern=PatternType.CONVENTION,  # Default, could be parsed from content
            action=action,
            title=title,
            content=content,
            evidence_session_ids=evidence_ids,
        )
        proposals.append(proposal)

    return proposals


@deal.pre(
    lambda session, max_tokens=4000: (
        isinstance(session, Session) and isinstance(max_tokens, int) and max_tokens > 0
    )
)
@deal.post(lambda result: isinstance(result, Session))
def truncate_session(
    session: Session,
    max_tokens: int = 4000,
) -> Session:
    """Truncate a session to fit within token budget.

    Keeps the most recent messages (from the end). Removes oldest messages
    first until total tokens <= max_tokens.

    Args:
        session: Session object to truncate.
        max_tokens: Maximum allowed token count.

    Returns:
        New Session with truncated entries and updated token_count.

    >>> from lattice.core.types.log import Session, LogEntry
    >>> from lattice.core.types.enums import Role
    >>> session = Session(session_id="test", logs=[])
    >>> truncated = truncate_session(session, max_tokens=100)
    >>> isinstance(truncated, Session)
    True
    """
    if not session.logs:
        return Session(session_id=session.session_id, logs=[])

    # Start from the end, build up logs
    kept_logs: list[LogEntry] = []
    current_tokens = 0

    for log in reversed(session.logs):
        # Rough token estimate: ~4 chars per token
        log_tokens = len(log.content) // 4 + 10  # +10 for metadata overhead

        if current_tokens + log_tokens <= max_tokens:
            kept_logs.insert(0, log)  # Insert at beginning to maintain order
            current_tokens += log_tokens
        else:
            break

    return Session(
        session_id=session.session_id,
        logs=kept_logs,
    )


@deal.pre(
    lambda sessions, batch_cap=30: isinstance(sessions, list)
    and isinstance(batch_cap, int)
)
@deal.post(lambda result: isinstance(result, list))
def apply_batch_cap(
    sessions: list[Session],
    batch_cap: int = 30,
) -> list[Session]:
    """Apply batch cap to limit sessions per evolution run.

    Takes the oldest sessions first (assumes sessions are already ordered
    by session_id, which corresponds to chronological order).

    Per RFC-002 §4.3: "If accumulated sessions exceed a configurable limit
    (default: 30 sessions), a single `evolve` processes the oldest batch
    up to the cap."

    Args:
        sessions: List of Session objects (should be ordered by session_id).
        batch_cap: Maximum number of sessions to return.

    Returns:
        List of at most `batch_cap` sessions (oldest first).

    >>> from lattice.core.types.log import Session
    >>> sessions = [Session(session_id=f"s{i}", logs=[]) for i in range(50)]
    >>> capped = apply_batch_cap(sessions, batch_cap=30)
    >>> len(capped)
    30
    >>> capped[0].session_id
    's0'
    >>> apply_batch_cap([], 30)
    []
    >>> apply_batch_cap(sessions, 0)
    []
    """
    if batch_cap <= 0:
        return []

    return sessions[:batch_cap]


# Re-export diff_rules from its own module for backward compatibility
from lattice.core.diff_rules import diff_rules  # noqa: E402
