"""LLM-based context compaction utility."""

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Dict, Optional

from emdash.agent.providers import get_provider
from emdash.agent.providers.base import LLMProvider


COMPACTION_SYSTEM_PROMPT = """You are a compaction engine.

You receive a JSON payload with sections that may be too large for context limits.
Return a single JSON object with:
- summary: short narrative of the payload (2-4 sentences)
- key_findings: 3-8 bullets
- open_questions: 0-5 bullets
- decisions: 0-5 bullets
- entities: list of key entities (names, file paths, identifiers)
- compacted_payload: same top-level keys as input payload, but shortened

Rules:
- Preserve identifiers, file paths, PR numbers, evidence IDs, and function/class names.
- Do NOT invent facts. If a section is empty, keep it empty.
- Output JSON only. No markdown fences or commentary.
"""

DEFAULT_TRIGGER_TOKENS = 12000
DEFAULT_TARGET_TOKENS = 6000


def _estimate_tokens(text: str) -> int:
    """Rough token estimate from character count."""
    if not text:
        return 0
    return max(1, len(text) // 4)


def _strip_json_fences(text: str) -> str:
    if "```" not in text:
        return text.strip()
    start = text.find("```")
    end = text.find("```", start + 3)
    if end > start:
        body = text[start + 3:end]
        if body.startswith("json"):
            body = body[4:]
        return body.strip()
    return text.strip()


@dataclass
class CompactionResult:
    summary: str
    key_findings: list[str]
    open_questions: list[str]
    decisions: list[str]
    entities: list[str]
    compacted_payload: Dict[str, str]

    def to_context_text(self) -> str:
        lines = ["COMPACTED CONTEXT", ""]
        if self.summary:
            lines.append(f"Summary: {self.summary}")
        if self.key_findings:
            lines.append("Key findings:")
            lines.extend([f"- {f}" for f in self.key_findings])
        if self.open_questions:
            lines.append("Open questions:")
            lines.extend([f"- {q}" for q in self.open_questions])
        if self.decisions:
            lines.append("Decisions:")
            lines.extend([f"- {d}" for d in self.decisions])
        if self.entities:
            lines.append("Entities:")
            lines.extend([f"- {e}" for e in self.entities])
        return "\n".join(lines).strip()


class LLMCompactor:
    """Reusable LLM-based compactor for large payloads."""

    def __init__(self, provider: Optional[LLMProvider] = None, model: Optional[str] = None):
        self.provider = provider or get_provider(model or "haiku")

    def compact_payload(
        self,
        payload: Dict[str, str],
        goal: str,
        trigger_ratio: float = 0.8,
        target_ratio: float = 0.35,
    ) -> Dict[str, str]:
        """Compact payload if it exceeds threshold.

        Returns the compacted payload (or original if no compaction needed).
        """
        payload_json = json.dumps(payload, ensure_ascii=True)
        context_limit = self.provider.get_context_limit()
        estimated_tokens = _estimate_tokens(payload_json)
        trigger_tokens = min(int(context_limit * trigger_ratio), DEFAULT_TRIGGER_TOKENS)
        if estimated_tokens < trigger_tokens:
            return payload

        target_tokens = min(int(context_limit * target_ratio), DEFAULT_TARGET_TOKENS)
        user_message = (
            f"GOAL: {goal}\n\n"
            f"TARGET_TOKENS: {target_tokens}\n\n"
            f"PAYLOAD JSON:\n{payload_json}"
        )
        messages = [{"role": "user", "content": user_message}]
        response = self.provider.chat(messages, system=COMPACTION_SYSTEM_PROMPT)
        content = response.content or ""

        try:
            data = json.loads(_strip_json_fences(content))
            compacted = data.get("compacted_payload")
            if not isinstance(compacted, dict):
                return payload
            # Ensure same keys exist
            return {k: compacted.get(k, v) for k, v in payload.items()}
        except Exception:
            # Fallback: truncate each section
            fallback = {}
            for key, value in payload.items():
                if not value:
                    fallback[key] = value
                    continue
                fallback[key] = value[:2000] + ("..." if len(value) > 2000 else "")
            return fallback
