"""Colony state management — YAML/JSONL-based persistence under .colony/."""

from __future__ import annotations

import json
from datetime import datetime, timezone
from pathlib import Path

import structlog
import yaml

from fliiq.runtime.colony.models import (
    AgentReflection,
    AgentTokenUsage,
    ColonyState,
    GovernanceDecision,
    IntelBrief,
    Proposal,
    ProposalStatus,
    QAMetrics,
    TimelineEntry,
    Vote,
)

log = structlog.get_logger()


def _normalize_role(data: dict, field: str) -> None:
    """Strip common LLM mistakes like '-agent' suffix from role fields."""
    val = data.get(field)
    if isinstance(val, str):
        data[field] = val.removesuffix("-agent").removesuffix("_agent")


class StateManager:
    """Manages colony state files under .colony/ directory."""

    def __init__(self, colony_dir: Path, state_subdir: str = "state"):
        self._dir = colony_dir
        self._state_dir = colony_dir / state_subdir

    def ensure_dirs(self) -> None:
        """Create the full directory structure."""
        dirs = [
            self._state_dir / "proposals",
            self._state_dir / "intelligence",
            self._state_dir / "intelligence" / "landscape",
            self._state_dir / "votes",
            self._dir / "reflections",
            self._state_dir / "metrics",
            self._dir / "reports",
        ]
        for d in dirs:
            d.mkdir(parents=True, exist_ok=True)

    def reset_working_state(self) -> None:
        """Clear proposals, votes, briefs, decisions, metrics, timeline.

        Preserves reflections (agent learning) and reports (historical).
        Used by directed mode to start with a clean slate.
        """
        import shutil

        # Clear file-per-item directories
        for subdir in ("proposals", "votes", "intelligence"):
            d = self._state_dir / subdir
            if d.is_dir():
                shutil.rmtree(d)

        # Clear append-only JSONL logs
        for name in (
            "governance_decisions.jsonl",
            "timeline.jsonl",
            "token_usage.jsonl",
            "mission_status.yaml",
        ):
            f = self._state_dir / name
            if f.is_file():
                f.unlink()

        # Clear metrics
        metrics_dir = self._state_dir / "metrics"
        if metrics_dir.is_dir():
            shutil.rmtree(metrics_dir)

        # Recreate empty directories
        self.ensure_dirs()
        log.info("colony.state_reset")

    # ── Proposals ──────────────────────────────────────────────────────

    def save_proposal(self, proposal: Proposal) -> Path:
        path = self._state_dir / "proposals" / f"{proposal.id}.yaml"
        path.write_text(yaml.dump(proposal.model_dump(mode="json"), default_flow_style=False, sort_keys=False))
        return path

    def load_proposal(self, proposal_id: str) -> Proposal | None:
        path = self._state_dir / "proposals" / f"{proposal_id}.yaml"
        if not path.is_file():
            return None
        data = yaml.safe_load(path.read_text())
        if not data:
            return None
        _normalize_role(data, "proposer")
        return Proposal(**data)

    def list_proposals(self, status: str | None = None) -> list[Proposal]:
        proposals_dir = self._state_dir / "proposals"
        if not proposals_dir.is_dir():
            return []
        proposals = []
        for f in sorted(proposals_dir.glob("PROP-*.yaml")):
            data = yaml.safe_load(f.read_text())
            if data:
                _normalize_role(data, "proposer")
                p = Proposal(**data)
                if status is None or p.status.value == status:
                    proposals.append(p)
        return proposals

    def update_proposal_status(self, proposal_id: str, status: str) -> None:
        p = self.load_proposal(proposal_id)
        if p:
            p.status = ProposalStatus(status)
            p.updated_at = datetime.now(timezone.utc)
            self.save_proposal(p)

    # ── Intelligence Briefs ────────────────────────────────────────────

    def save_intel_brief(self, brief: IntelBrief) -> Path:
        path = self._state_dir / "intelligence" / f"{brief.id}.yaml"
        path.write_text(yaml.dump(brief.model_dump(mode="json"), default_flow_style=False, sort_keys=False))
        return path

    def load_intel_brief(self, brief_id: str) -> IntelBrief | None:
        path = self._state_dir / "intelligence" / f"{brief_id}.yaml"
        if not path.is_file():
            return None
        data = yaml.safe_load(path.read_text())
        return IntelBrief(**data) if data else None

    def list_intel_briefs(self, since: datetime | None = None) -> list[IntelBrief]:
        intel_dir = self._state_dir / "intelligence"
        if not intel_dir.is_dir():
            return []
        briefs = []
        for f in sorted(intel_dir.glob("INTEL-*.yaml")):
            data = yaml.safe_load(f.read_text())
            if data:
                b = IntelBrief(**data)
                if since is None or b.created_at >= since:
                    briefs.append(b)
        return briefs

    # ── Votes ──────────────────────────────────────────────────────────

    def save_vote(self, vote: Vote) -> Path:
        path = self._state_dir / "votes" / f"{vote.proposal_id}_{vote.voter.value}.yaml"
        path.write_text(yaml.dump(vote.model_dump(mode="json"), default_flow_style=False, sort_keys=False))
        return path

    def get_votes_for(self, proposal_id: str) -> list[Vote]:
        votes_dir = self._state_dir / "votes"
        if not votes_dir.is_dir():
            return []
        votes = []
        for f in votes_dir.glob(f"{proposal_id}_*.yaml"):
            data = yaml.safe_load(f.read_text())
            if data:
                _normalize_role(data, "voter")
                votes.append(Vote(**data))
        return votes

    # ── Reflections ────────────────────────────────────────────────────

    def append_reflection(self, reflection: AgentReflection) -> None:
        path = self._dir / "reflections" / f"{reflection.agent.value}.jsonl"
        with path.open("a") as f:
            f.write(json.dumps(reflection.model_dump(mode="json")) + "\n")

    def get_reflections(self, agent: str, limit: int = 5) -> list[AgentReflection]:
        path = self._dir / "reflections" / f"{agent}.jsonl"
        if not path.is_file():
            return []
        lines = path.read_text().strip().splitlines()
        entries = []
        for idx, line in enumerate(lines[-limit:]):
            data = json.loads(line)
            # Backfill required fields if agent wrote them without
            data.setdefault("agent", agent)
            data.setdefault("cycle_number", idx)
            entries.append(AgentReflection(**data))
        return entries

    # ── Token Usage ─────────────────────────────────────────────────────

    def append_token_usage(self, usage: AgentTokenUsage) -> None:
        path = self._state_dir / "token_usage.jsonl"
        with path.open("a") as f:
            f.write(json.dumps(usage.model_dump(mode="json")) + "\n")

    def get_token_usage(self) -> list[AgentTokenUsage]:
        path = self._state_dir / "token_usage.jsonl"
        if not path.is_file():
            return []
        entries = []
        for line in path.read_text().strip().splitlines():
            entries.append(AgentTokenUsage(**json.loads(line)))
        return entries

    # ── Governance Decisions ───────────────────────────────────────────

    def append_decision(self, decision: GovernanceDecision) -> None:
        path = self._state_dir / "governance_decisions.jsonl"
        with path.open("a") as f:
            f.write(json.dumps(decision.model_dump(mode="json")) + "\n")

    def get_decisions(self) -> list[GovernanceDecision]:
        path = self._state_dir / "governance_decisions.jsonl"
        if not path.is_file():
            return []
        decisions = []
        for line in path.read_text().strip().splitlines():
            decisions.append(GovernanceDecision(**json.loads(line)))
        return decisions

    # ── QA Metrics ─────────────────────────────────────────────────────

    def save_metrics(self, metrics: QAMetrics) -> None:
        path = self._state_dir / "metrics" / "metrics.jsonl"
        with path.open("a") as f:
            f.write(json.dumps(metrics.model_dump(mode="json")) + "\n")

    def get_latest_metrics(self) -> QAMetrics | None:
        path = self._state_dir / "metrics" / "metrics.jsonl"
        if not path.is_file():
            return None
        lines = path.read_text().strip().splitlines()
        if not lines:
            return None
        return QAMetrics(**json.loads(lines[-1]))

    def get_metrics_history(self) -> list[QAMetrics]:
        path = self._state_dir / "metrics" / "metrics.jsonl"
        if not path.is_file():
            return []
        return [QAMetrics(**json.loads(line)) for line in path.read_text().strip().splitlines()]

    # ── Timeline ───────────────────────────────────────────────────────

    def append_timeline(self, entry: TimelineEntry) -> None:
        path = self._state_dir / "timeline.jsonl"
        with path.open("a") as f:
            f.write(json.dumps(entry.model_dump(mode="json")) + "\n")

    def get_timeline(self, limit: int = 0) -> list[TimelineEntry]:
        path = self._state_dir / "timeline.jsonl"
        if not path.is_file():
            return []
        lines = path.read_text().strip().splitlines()
        if limit:
            lines = lines[-limit:]
        return [TimelineEntry(**json.loads(line)) for line in lines]

    # ── Colony State ───────────────────────────────────────────────────

    def save_colony_state(self, state: ColonyState) -> None:
        path = self._state_dir / "colony_state.yaml"
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(yaml.dump(state.model_dump(mode="json"), default_flow_style=False, sort_keys=False))

    def load_colony_state(self) -> ColonyState | None:
        path = self._state_dir / "colony_state.yaml"
        if not path.is_file():
            return None
        data = yaml.safe_load(path.read_text())
        return ColonyState(**data) if data else None

    # ── Context Builders (for agent prompts) ───────────────────────────

    def build_intelligence_context(self, limit: int = 10) -> str:
        """Render recent intelligence briefs as text for agent prompts."""
        briefs = self.list_intel_briefs()[-limit:]
        if not briefs:
            return "No intelligence briefs yet."
        lines = []
        for b in briefs:
            lines.append(f"### {b.id}: {b.title}")
            lines.append(f"Relevance: {b.relevance} | Category: {b.category} | Source: {b.source_type}")
            lines.append(b.summary)
            if b.implications:
                lines.append("Implications: " + "; ".join(b.implications))
            lines.append("")
        return "\n".join(lines)

    def build_proposals_context(self, status: str | None = None) -> str:
        """Render proposals as text for agent prompts."""
        proposals = self.list_proposals(status=status)
        if not proposals:
            return "No proposals yet." if status is None else f"No proposals with status '{status}'."
        lines = []
        for p in proposals:
            lines.append(f"### {p.id}: {p.title}")
            lines.append(f"Type: {p.type} | Risk: {p.risk} | Status: {p.status} | By: {p.proposer}")
            lines.append(p.description[:500])
            votes = self.get_votes_for(p.id)
            if votes:
                for v in votes:
                    lines.append(f"  Vote by {v.voter}: {v.position} — {v.reasoning[:200]}")
            lines.append("")
        return "\n".join(lines)

    def build_qa_context(self) -> str:
        """Render latest metrics + recent history as text."""
        latest = self.get_latest_metrics()
        if not latest:
            return "No QA metrics recorded yet."
        history = self.get_metrics_history()[-5:]
        lines = [
            f"Latest: {latest.pass_count}/{latest.test_count} tests passing, {latest.lint_issues} lint issues",
        ]
        if len(history) > 1:
            lines.append("Recent trend:")
            for m in history:
                lines.append(f"  {m.timestamp}: {m.pass_count}/{m.test_count} pass, {m.fail_count} fail")
        return "\n".join(lines)
