"""Read/write context files, branch dir management."""

from __future__ import annotations

import json
import os
import shutil
from pathlib import Path
from typing import Optional

from pjctx.core.config import get_pjctx_dir, load_config
from pjctx.core.context import Context

CONTEXTS_DIR = "contexts"


def _branch_dir(repo_root: Path, branch: str) -> Path:
    """Get context directory for a branch (nested paths preserved)."""
    return get_pjctx_dir(repo_root) / CONTEXTS_DIR / branch


def save_context(repo_root: Path, ctx: Context) -> Path:
    """Save context to timestamped file and update latest.json. Returns path."""
    branch_dir = _branch_dir(repo_root, ctx.branch)
    branch_dir.mkdir(parents=True, exist_ok=True)

    data = ctx.to_dict()
    json_bytes = json.dumps(data, indent=2) + "\n"

    # Write timestamped file (atomic)
    ts_path = branch_dir / ctx.filename()
    tmp_path = ts_path.with_suffix(".tmp")
    tmp_path.write_text(json_bytes)
    os.replace(tmp_path, ts_path)

    # Write latest.json (file copy, not symlink)
    latest_path = branch_dir / "latest.json"
    tmp_latest = latest_path.with_suffix(".tmp")
    tmp_latest.write_text(json_bytes)
    os.replace(tmp_latest, latest_path)

    # Prune old contexts
    _prune(repo_root, ctx.branch)

    return ts_path


def load_latest(repo_root: Path, branch: str) -> Optional[Context]:
    """Load the most recent context for a branch."""
    latest = _branch_dir(repo_root, branch) / "latest.json"
    if not latest.exists():
        return None
    data = json.loads(latest.read_text())
    return Context.from_dict(data)


def list_contexts(repo_root: Path, branch: str) -> list[Context]:
    """List all contexts for a branch, sorted oldest first."""
    branch_dir = _branch_dir(repo_root, branch)
    if not branch_dir.exists():
        return []
    contexts: list[Context] = []
    for f in sorted(branch_dir.glob("*.json")):
        if f.name == "latest.json":
            continue
        try:
            data = json.loads(f.read_text())
            contexts.append(Context.from_dict(data))
        except (json.JSONDecodeError, KeyError):
            continue
    return contexts


def list_branches(repo_root: Path) -> list[str]:
    """List all branches that have saved contexts."""
    ctx_dir = get_pjctx_dir(repo_root) / CONTEXTS_DIR
    if not ctx_dir.exists():
        return []
    branches: list[str] = []
    for path in ctx_dir.rglob("latest.json"):
        rel = path.parent.relative_to(ctx_dir)
        branches.append(str(rel))
    return sorted(branches)


def _prune(repo_root: Path, branch: str) -> None:
    """Remove oldest contexts if exceeding max_contexts_per_branch."""
    config = load_config(repo_root)
    max_count = config.get("max_contexts_per_branch", 50)
    branch_dir = _branch_dir(repo_root, branch)
    files = sorted(
        [f for f in branch_dir.glob("*.json") if f.name != "latest.json"]
    )
    while len(files) > max_count:
        files[0].unlink()
        files.pop(0)
