"""Experiment and backlog persistence — read/write TOML files, manage indexes."""

from __future__ import annotations

import json
import tomllib
from pathlib import Path

import tomli_w

from abreka.config import load_index, save_index
from abreka.errors import BacklogItemNotFoundError, ExperimentNotFoundError
from abreka.models import BacklogItem, Experiment, ExperimentStatus

try:
    import fcntl
except ImportError:  # pragma: no cover - Windows fallback path
    fcntl = None


def experiments_dir(abreka_dir: Path) -> Path:
    return abreka_dir / "experiments"


def experiment_dir(abreka_dir: Path, exp_id: str) -> Path:
    return experiments_dir(abreka_dir) / exp_id


def next_experiment_id(abreka_dir: Path) -> str:
    """Scan existing experiment directories and return the next zero-padded 4-digit ID."""
    exp_dir = experiments_dir(abreka_dir)
    if not exp_dir.is_dir():
        return "0001"
    existing = sorted(
        d.name for d in exp_dir.iterdir()
        if d.is_dir() and d.name.isdigit() and len(d.name) == 4
    )
    if not existing:
        return "0001"
    return f"{int(existing[-1]) + 1:04d}"


def normalize_exp_id(raw: str) -> str:
    """Normalize an experiment ID to zero-padded 4-digit format."""
    return f"{int(raw):04d}"


def allocate_experiment_id(abreka_dir: Path) -> str:
    """Allocate the next experiment id with a lock-backed counter."""
    counter_path = abreka_dir / ".next_experiment_id"
    lock_path = abreka_dir / ".next_experiment_id.lock"
    lock_path.parent.mkdir(parents=True, exist_ok=True)

    with open(lock_path, "a+") as lock_file:
        if fcntl is not None:
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
        scanned_next = int(next_experiment_id(abreka_dir))
        counter_next = scanned_next
        if counter_path.is_file():
            try:
                counter_next = int(counter_path.read_text().strip())
            except ValueError:
                counter_next = scanned_next
        next_id = max(scanned_next, counter_next)
        counter_path.write_text(f"{next_id + 1:04d}\n")
        return f"{next_id:04d}"


def reserve_experiment(abreka_dir: Path) -> Experiment:
    """Reserve an experiment id by creating a stub experiment in proposing state."""
    exp = Experiment(
        id=allocate_experiment_id(abreka_dir),
        status=ExperimentStatus.PROPOSING,
    )
    create_experiment_dir(abreka_dir, exp.id)
    save_experiment(exp, abreka_dir)
    update_index_entry(abreka_dir, exp)
    return exp


def create_experiment_dir(abreka_dir: Path, exp_id: str) -> Path:
    """Create the experiment directory structure."""
    exp = experiment_dir(abreka_dir, exp_id)
    exp.mkdir(parents=True, exist_ok=True)
    (exp / "code").mkdir(exist_ok=True)
    (exp / "artifacts" / "checkpoints").mkdir(parents=True, exist_ok=True)
    (exp / "artifacts" / "plots").mkdir(parents=True, exist_ok=True)
    (exp / "artifacts" / "eda").mkdir(parents=True, exist_ok=True)
    return exp


def load_experiment(abreka_dir: Path, exp_id: str) -> Experiment:
    """Load an experiment from its experiment.toml file."""
    toml_path = experiment_dir(abreka_dir, exp_id) / "experiment.toml"
    if not toml_path.is_file():
        raise ExperimentNotFoundError(f"Experiment {exp_id} not found: {toml_path}")
    with open(toml_path, "rb") as f:
        data = tomllib.load(f)
    return Experiment.from_dict(data)


def save_experiment(exp: Experiment, abreka_dir: Path) -> None:
    """Write an experiment to its experiment.toml file."""
    toml_path = experiment_dir(abreka_dir, exp.id) / "experiment.toml"
    toml_path.parent.mkdir(parents=True, exist_ok=True)
    with open(toml_path, "wb") as f:
        tomli_w.dump(exp.to_dict(), f)


def update_index_entry(abreka_dir: Path, exp: Experiment) -> None:
    """Update a single experiment's entry in index.json."""
    index = load_index(abreka_dir)
    index[exp.id] = exp.summary()
    save_index(index, abreka_dir)


def rebuild_index(abreka_dir: Path) -> dict:
    """Rebuild the entire index by scanning all experiment.toml files."""
    index: dict = {}
    exp_root = experiments_dir(abreka_dir)
    if not exp_root.is_dir():
        save_index(index, abreka_dir)
        return index
    for exp_dir in sorted(exp_root.iterdir()):
        toml_path = exp_dir / "experiment.toml"
        if exp_dir.is_dir() and toml_path.is_file():
            with open(toml_path, "rb") as f:
                data = tomllib.load(f)
            exp = Experiment.from_dict(data)
            index[exp.id] = exp.summary()
    save_index(index, abreka_dir)
    return index


# -- Backlog persistence -------------------------------------------------------


def backlog_dir(abreka_dir: Path) -> Path:
    return abreka_dir / "backlog"


def next_backlog_id(abreka_dir: Path) -> str:
    """Scan existing backlog TOML files and return the next zero-padded 4-digit ID."""
    bl_dir = backlog_dir(abreka_dir)
    if not bl_dir.is_dir():
        return "0001"
    existing = sorted(
        f.stem for f in bl_dir.iterdir()
        if f.is_file() and f.suffix == ".toml" and f.stem.isdigit() and len(f.stem) == 4
    )
    if not existing:
        return "0001"
    return f"{int(existing[-1]) + 1:04d}"


def normalize_backlog_id(raw: str) -> str:
    """Normalize a backlog item ID to zero-padded 4-digit format.

    Accepts a leading ``B``/``b`` prefix because branch summaries render backlog
    ids that way (e.g. ``B0002``), so agents frequently pass the prefixed form.
    """
    stripped = raw.lstrip("Bb") if raw else raw
    return f"{int(stripped):04d}"


def load_backlog_item(abreka_dir: Path, item_id: str) -> BacklogItem:
    """Load a backlog item from its TOML file."""
    toml_path = backlog_dir(abreka_dir) / f"{item_id}.toml"
    if not toml_path.is_file():
        raise BacklogItemNotFoundError(f"Backlog item {item_id} not found: {toml_path}")
    with open(toml_path, "rb") as f:
        data = tomllib.load(f)
    return BacklogItem.from_dict(data)


def save_backlog_item(item: BacklogItem, abreka_dir: Path) -> None:
    """Write a backlog item to its TOML file."""
    bl_dir = backlog_dir(abreka_dir)
    bl_dir.mkdir(parents=True, exist_ok=True)
    toml_path = bl_dir / f"{item.id}.toml"
    with open(toml_path, "wb") as f:
        tomli_w.dump(item.to_dict(), f)


def delete_backlog_item(abreka_dir: Path, item_id: str) -> None:
    """Remove a backlog item's TOML file and index entry."""
    toml_path = backlog_dir(abreka_dir) / f"{item_id}.toml"
    if toml_path.is_file():
        toml_path.unlink()
    index = load_backlog_index(abreka_dir)
    index.pop(item_id, None)
    save_backlog_index(index, abreka_dir)


def load_backlog_index(abreka_dir: Path) -> dict:
    """Read abreka/backlog_index.json."""
    index_path = abreka_dir / "backlog_index.json"
    if index_path.is_file():
        with open(index_path) as f:
            return json.load(f)
    return {}


def save_backlog_index(index: dict, abreka_dir: Path) -> None:
    """Write abreka/backlog_index.json."""
    index_path = abreka_dir / "backlog_index.json"
    with open(index_path, "w") as f:
        json.dump(index, f, indent=2)
        f.write("\n")


def update_backlog_index_entry(abreka_dir: Path, item: BacklogItem) -> None:
    """Update a single backlog item's entry in backlog_index.json."""
    index = load_backlog_index(abreka_dir)
    index[item.id] = item.summary()
    save_backlog_index(index, abreka_dir)


def rebuild_backlog_index(abreka_dir: Path) -> dict:
    """Rebuild the entire backlog index by scanning all backlog/*.toml files."""
    index: dict = {}
    bl_dir = backlog_dir(abreka_dir)
    if not bl_dir.is_dir():
        save_backlog_index(index, abreka_dir)
        return index
    for toml_path in sorted(bl_dir.glob("*.toml")):
        with open(toml_path, "rb") as f:
            data = tomllib.load(f)
        item = BacklogItem.from_dict(data)
        index[item.id] = item.summary()
    save_backlog_index(index, abreka_dir)
    return index


