from __future__ import annotations

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

from pydantic import BaseModel

from epochler.app_settings import get_workspace_root
from epochler.config import get_settings
from epochler.contract.models import (
    AgentActivity,
    ContractPhase,
    Conversation,
    ExperimentJournal,
    ExperimentResult,
    FrozenContract,
    PhaseArtifact,
)

T = TypeVar("T", bound=BaseModel)
_EXPERIMENT_CODE_FILENAME = "train_code.txt"


def _data_dir() -> Path:
    d = get_settings().data_dir
    d.mkdir(parents=True, exist_ok=True)
    return d


def _ensure_dir(path: Path) -> Path:
    path.mkdir(parents=True, exist_ok=True)
    return path


def _write_json(path: Path, model: BaseModel) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(model.model_dump_json(indent=2), encoding="utf-8")


def _read_json(path: Path, cls: type[T]) -> T | None:
    if not path.exists():
        return None
    return cls.model_validate_json(path.read_text(encoding="utf-8"))


# ---------------------------------------------------------------------------
# Conversations
# ---------------------------------------------------------------------------

def conversations_dir() -> Path:
    return _ensure_dir(_data_dir() / "conversations")


def save_conversation(conv: Conversation) -> None:
    conv.updated_at = datetime.now(timezone.utc)
    _write_json(conversations_dir() / f"{conv.id}.json", conv)


def load_conversation(conv_id: str) -> Conversation | None:
    return _read_json(conversations_dir() / f"{conv_id}.json", Conversation)


def list_conversations() -> list[Conversation]:
    results = []
    for p in sorted(conversations_dir().glob("*.json"), key=lambda f: f.stat().st_mtime, reverse=True):
        conv = _read_json(p, Conversation)
        if conv:
            results.append(conv)
    return results


def delete_conversation(conv_id: str, delete_files: bool = False) -> bool:
    """Delete a conversation and optionally its associated workspace/experiment files."""
    import shutil

    conv = load_conversation(conv_id)
    path = conversations_dir() / f"{conv_id}.json"
    if not path.exists():
        return False

    contract_id = conv.contract.id if conv else None

    path.unlink()

    if delete_files and contract_id:
        # Delete contract JSON
        contract_path = contracts_dir() / f"{contract_id}.json"
        if contract_path.exists():
            contract_path.unlink()

        # Delete experiments dir
        exp_dir = _data_dir() / "experiments" / contract_id
        if exp_dir.exists():
            shutil.rmtree(exp_dir, ignore_errors=True)

        # Delete workspace dir (the generated execution workspace)
        ws_root = conv.workspace_root if conv else None
        ws = workspaces_dir(ws_root or None) / contract_id
        if ws.exists():
            shutil.rmtree(ws, ignore_errors=True)

    return True


# ---------------------------------------------------------------------------
# Contracts
# ---------------------------------------------------------------------------

def contracts_dir() -> Path:
    return _ensure_dir(_data_dir() / "contracts")


def save_contract(contract: FrozenContract) -> None:
    _write_json(contracts_dir() / f"{contract.id}.json", contract)


def load_contract(contract_id: str) -> FrozenContract | None:
    return _read_json(contracts_dir() / f"{contract_id}.json", FrozenContract)


# ---------------------------------------------------------------------------
# Experiments
# ---------------------------------------------------------------------------

def experiments_dir(contract_id: str) -> Path:
    return _ensure_dir(_data_dir() / "experiments" / contract_id)


def save_experiment(result: ExperimentResult) -> None:
    run_dir = _ensure_dir(experiments_dir(result.contract_id) / result.id)
    _write_json(run_dir / "journal.json", result)


def load_experiment(contract_id: str, experiment_id: str) -> ExperimentResult | None:
    return _read_json(
        experiments_dir(contract_id) / experiment_id / "journal.json",
        ExperimentResult,
    )


def list_experiments(contract_id: str) -> list[ExperimentResult]:
    base = experiments_dir(contract_id)
    results = []
    for run_dir in sorted(base.iterdir()):
        journal = run_dir / "journal.json"
        if journal.exists():
            exp = _read_json(journal, ExperimentResult)
            if exp:
                results.append(exp)
    return results


def save_experiment_code(contract_id: str, experiment_id: str, code: str) -> None:
    run_dir = _ensure_dir(experiments_dir(contract_id) / experiment_id)
    # Store archived code under a non-.py filename so dev autoreload does not
    # mistake experiment snapshots for backend source edits mid-run.
    (run_dir / _EXPERIMENT_CODE_FILENAME).write_text(code, encoding="utf-8")


def load_experiment_code(contract_id: str, experiment_id: str) -> str | None:
    path = experiments_dir(contract_id) / experiment_id / _EXPERIMENT_CODE_FILENAME
    return path.read_text(encoding="utf-8") if path.exists() else None


def save_experiment_logs(
    contract_id: str, experiment_id: str, stdout: str, stderr: str
) -> None:
    run_dir = _ensure_dir(experiments_dir(contract_id) / experiment_id)
    (run_dir / "stdout.log").write_text(stdout, encoding="utf-8")
    (run_dir / "stderr.log").write_text(stderr, encoding="utf-8")


def save_experiment_phase_outputs(
    contract_id: str,
    experiment_id: str,
    phase_outputs: dict[str, dict[str, str]],
) -> None:
    run_dir = _ensure_dir(experiments_dir(contract_id) / experiment_id / "phases")
    for phase_name, streams in phase_outputs.items():
        phase_dir = _ensure_dir(run_dir / phase_name)
        for filename, content in streams.items():
            (phase_dir / filename).write_text(content, encoding="utf-8")


# ---------------------------------------------------------------------------
# Experiment journal (aggregate)
# ---------------------------------------------------------------------------

def save_journal(journal: ExperimentJournal) -> None:
    _write_json(experiments_dir(journal.contract_id) / "journal.json", journal)


def load_journal(contract_id: str) -> ExperimentJournal | None:
    return _read_json(experiments_dir(contract_id) / "journal.json", ExperimentJournal)


def clear_journal(contract_id: str) -> None:
    path = experiments_dir(contract_id) / "journal.json"
    if path.exists():
        path.unlink()


# ---------------------------------------------------------------------------
# Activity log (append-only JSONL)
# ---------------------------------------------------------------------------

def append_activity(contract_id: str, activity: AgentActivity) -> None:
    path = experiments_dir(contract_id) / "activity_log.jsonl"
    with open(path, "a", encoding="utf-8") as f:
        f.write(activity.model_dump_json() + "\n")


def load_activity_log(contract_id: str) -> list[AgentActivity]:
    path = experiments_dir(contract_id) / "activity_log.jsonl"
    if not path.exists():
        return []
    activities = []
    for line in path.read_text(encoding="utf-8").strip().splitlines():
        if line.strip():
            activities.append(AgentActivity.model_validate_json(line))
    return activities


# ---------------------------------------------------------------------------
# Phase manifests and artifacts
# ---------------------------------------------------------------------------

def phases_dir(contract_id: str) -> Path:
    return _ensure_dir(experiments_dir(contract_id) / "phases")


def phase_dir(contract_id: str, phase_name: str) -> Path:
    return _ensure_dir(phases_dir(contract_id) / phase_name)


def save_phase_manifest(contract_id: str, phase: ContractPhase) -> None:
    _write_json(phase_dir(contract_id, phase.name.value) / "manifest.json", phase)


def load_phase_manifest(contract_id: str, phase_name: str) -> ContractPhase | None:
    return _read_json(phase_dir(contract_id, phase_name) / "manifest.json", ContractPhase)


def list_phase_manifests(contract_id: str) -> dict[str, ContractPhase]:
    manifests: dict[str, ContractPhase] = {}
    for phase_path in sorted(phases_dir(contract_id).iterdir()):
        if not phase_path.is_dir():
            continue
        manifest = _read_json(phase_path / "manifest.json", ContractPhase)
        if manifest:
            manifests[phase_path.name] = manifest
    return manifests


def save_phase_script(contract_id: str, phase_name: str, code: str) -> Path:
    path = workspace_path(contract_id) / f"{phase_name}.py"
    path.write_text(code, encoding="utf-8")
    return path


def load_phase_script(contract_id: str, phase_name: str) -> str | None:
    path = workspace_path(contract_id) / f"{phase_name}.py"
    return path.read_text(encoding="utf-8") if path.exists() else None


def save_phase_artifacts(contract_id: str, phase_name: str, artifacts: list[PhaseArtifact]) -> None:
    payload = {"phase_name": phase_name, "artifacts": [artifact.model_dump(mode="json") for artifact in artifacts]}
    path = phase_dir(contract_id, phase_name) / "artifacts.json"
    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")


def load_phase_artifacts(contract_id: str, phase_name: str) -> list[PhaseArtifact]:
    path = phase_dir(contract_id, phase_name) / "artifacts.json"
    if not path.exists():
        return []
    data = json.loads(path.read_text(encoding="utf-8"))
    return [PhaseArtifact(**artifact) for artifact in data.get("artifacts", [])]


def artifact_versions_dir(contract_id: str, phase_name: str) -> Path:
    return _ensure_dir(phases_dir(contract_id) / phase_name / "artifact_versions")


def save_artifact_version(contract_id: str, phase_name: str, artifact: PhaseArtifact) -> Path:
    fingerprint = artifact.fingerprint or "unversioned"
    path = artifact_versions_dir(contract_id, phase_name) / f"{fingerprint}.json"
    path.write_text(json.dumps(artifact.model_dump(mode="json"), indent=2), encoding="utf-8")
    return path


def list_artifact_versions(contract_id: str, phase_name: str) -> list[PhaseArtifact]:
    versions: list[PhaseArtifact] = []
    for version_path in sorted(artifact_versions_dir(contract_id, phase_name).glob("*.json")):
        data = json.loads(version_path.read_text(encoding="utf-8"))
        versions.append(PhaseArtifact(**data))
    return versions


# ---------------------------------------------------------------------------
# Workspaces
# ---------------------------------------------------------------------------

def workspaces_dir(override_root: str | None = None) -> Path:
    if override_root:
        return _ensure_dir(Path(override_root))
    root = get_workspace_root()
    if root:
        return _ensure_dir(root)
    return _ensure_dir(_data_dir() / "workspaces")


def workspace_path(contract_id: str, override_root: str | None = None) -> Path:
    return _ensure_dir(workspaces_dir(override_root) / contract_id)
