"""Policy CRUD endpoints for MemoryPolicy objects."""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timedelta
import os
from pathlib import Path
from typing import Any

from fastapi import APIRouter, HTTPException, Query
import networkx as nx
import yaml

from hieromem_api.models.requests import PolicySimulationRequest, PolicyUpsertRequest
from hieromem_api.models.responses import (
    PolicyDecisionResponse,
    PolicyDocument,
    PolicyListResponse,
    PolicySimulationResponse,
    PolicyUpsertResponse,
)
from hieromem_core.policy.parser import load_policy
from hieromem_core.policy.schema import MemoryPolicy
from hieromem_core.policy.evaluator import MemoryPolicyEvaluator

router = APIRouter(prefix="/policies", tags=["Policies"])


@dataclass
class StoredPolicy:
    policy: MemoryPolicy
    raw: dict[str, Any]


_POLICY_STORE: dict[tuple[str, str], StoredPolicy] = {}


def _policy_key(namespace: str | None, name: str) -> tuple[str, str]:
    return (namespace or "default", name)


def _load_default_policies() -> None:
    default_dir = Path.cwd() / "configs" / "policies"
    policy_dir = Path(os.environ.get("HIEROMEM_POLICY_DIR", default_dir))
    if not policy_dir.exists():
        return
    for path in sorted(policy_dir.glob("*.yaml")):
        payload = yaml.safe_load(path.read_text(encoding="utf-8"))
        if not isinstance(payload, dict):
            continue
        policy = load_policy(payload)
        namespace = policy.namespace or "default"
        _POLICY_STORE[_policy_key(namespace, policy.name)] = StoredPolicy(
            policy=policy,
            raw=policy.to_dict(),
        )


@router.get("", response_model=PolicyListResponse)
def list_policies() -> PolicyListResponse:
    documents = [
        PolicyDocument(
            name=stored.policy.name,
            namespace=stored.policy.namespace,
            policy=stored.raw,
        )
        for stored in _POLICY_STORE.values()
    ]
    return PolicyListResponse(policies=documents)


@router.get("/{name}", response_model=PolicyDocument)
def get_policy(
    name: str,
    namespace: str | None = Query(default=None),
) -> PolicyDocument:
    key = _policy_key(namespace, name)
    stored = _POLICY_STORE.get(key)
    if stored is None:
        raise HTTPException(status_code=404, detail="Policy not found")
    return PolicyDocument(
        name=stored.policy.name,
        namespace=stored.policy.namespace,
        policy=stored.raw,
    )


@router.post("", response_model=PolicyUpsertResponse)
def upsert_policy(request: PolicyUpsertRequest) -> PolicyUpsertResponse:
    policy = load_policy(request.policy)
    namespace = policy.namespace or "default"
    document = PolicyDocument(
        name=policy.name,
        namespace=policy.namespace,
        policy=policy.to_dict(),
    )
    if not request.validate_only:
        _POLICY_STORE[_policy_key(namespace, policy.name)] = StoredPolicy(
            policy=policy,
            raw=policy.to_dict(),
        )
        status = "saved"
    else:
        status = "validated"
    return PolicyUpsertResponse(status=status, document=document)


@router.post("/simulate", response_model=PolicySimulationResponse)
def simulate_policy(request: PolicySimulationRequest) -> PolicySimulationResponse:
    policy = load_policy(request.policy)
    evaluator = MemoryPolicyEvaluator(policy)
    namespace = request.namespace or policy.namespace or "default"

    @dataclass
    class _EpisodeStub:
        id: str
        tocc: datetime

    graph = None
    if request.graph_edges:
        graph = nx.DiGraph()
        for edge in request.graph_edges:
            attrs: dict[str, Any] = {}
            if edge.relation:
                attrs["relation"] = edge.relation
            graph.add_edge(edge.source, edge.target, **attrs)

    insert_decision = evaluator.evaluate_insert(
        graph=graph,
        total_episodes=request.total_episodes,
        namespace=namespace,
    )

    episode_count = request.episode_count or 0
    oldest_age = request.oldest_episode_age_s or 0
    now = datetime.now()
    episodes = [
        _EpisodeStub(id=f"episode_{idx}", tocc=now)
        for idx in range(episode_count)
    ]
    if episodes and oldest_age:
        episodes[0].tocc = now - timedelta(seconds=oldest_age)

    summarise_decision = evaluator.evaluate_summarise(episodes, namespace=namespace)
    retention_candidates = [
        ep.id for ep in evaluator.retention_candidates(episodes)
    ]

    return PolicySimulationResponse(
        namespace=namespace,
        insert=PolicyDecisionResponse(**insert_decision.__dict__),
        summarise=PolicyDecisionResponse(**summarise_decision.__dict__),
        retention_candidates=retention_candidates,
    )


@router.delete("/{name}", response_model=PolicyUpsertResponse)
def delete_policy(
    name: str,
    namespace: str | None = Query(default=None),
) -> PolicyUpsertResponse:
    key = _policy_key(namespace, name)
    stored = _POLICY_STORE.pop(key, None)
    if stored is None:
        raise HTTPException(status_code=404, detail="Policy not found")
    document = PolicyDocument(
        name=stored.policy.name,
        namespace=stored.policy.namespace,
        policy=stored.raw,
    )
    return PolicyUpsertResponse(status="deleted", document=document)


_load_default_policies()
