"""abreka exp chat — launch an interactive pi session briefed on one experiment."""

from __future__ import annotations

import os
import shutil
from importlib.resources import files as pkg_files
from pathlib import Path

import click

from abreka.config import discover_project, find_abreka_dir, load_config, load_goal, load_index
from abreka.models import Experiment
from abreka.persistence import experiment_dir, load_experiment, normalize_exp_id
from abreka.prompts import render_template


_TEMPLATE_RESOURCE = pkg_files("abreka.defaults") / "exp_chat_prompt.md"


def _format_metrics_table(metrics: dict[str, dict[str, float]]) -> str:
    if not metrics:
        return "(no metrics recorded)"
    lines = []
    for split in sorted(metrics):
        for name in sorted(metrics[split]):
            value = metrics[split][name]
            lines.append(f"  {split}:{name} = {value:.6f}")
    return "\n".join(lines)


def _format_parent_summaries(abreka_dir: Path, parent_ids: list[str]) -> str:
    if not parent_ids:
        return "(no parents — this is a root experiment)"
    index = load_index(abreka_dir)
    lines = []
    for pid in parent_ids:
        entry = index.get(pid)
        if not entry:
            lines.append(f"- {pid}: (not found in index)")
            continue
        status = entry.get("status", "?")
        hypothesis = entry.get("hypothesis", "").strip() or "(no hypothesis)"
        lines.append(f"- {pid} [{status}]: {hypothesis}")
    return "\n".join(lines)


def _list_experiment_files(exp_path: Path) -> str:
    if not exp_path.is_dir():
        return "(experiment directory missing)"
    entries: list[str] = []
    for name in ("run.py", "findings.md", "experiment.toml"):
        p = exp_path / name
        marker = "present" if p.is_file() else "missing"
        entries.append(f"- {name}: {marker}")
    code_dir = exp_path / "code"
    if code_dir.is_dir():
        code_files = sorted(p.name for p in code_dir.iterdir() if p.is_file())
        entries.append(f"- code/: {', '.join(code_files) if code_files else '(empty)'}")
    artifacts = exp_path / "artifacts"
    if artifacts.is_dir():
        checkpoints = artifacts / "checkpoints"
        plots = artifacts / "plots"
        if checkpoints.is_dir():
            n = sum(1 for _ in checkpoints.iterdir())
            entries.append(f"- artifacts/checkpoints/: {n} file(s)")
        if plots.is_dir():
            n = sum(1 for _ in plots.iterdir())
            entries.append(f"- artifacts/plots/: {n} file(s)")
    return "\n".join(entries)


def _build_context(abreka_dir: Path, exp: Experiment) -> dict[str, str]:
    project_root = abreka_dir.parent
    project_name, lib_name = discover_project(project_root)
    config = load_config(abreka_dir)
    goal_text = load_goal(abreka_dir).strip() or "(no goal.md written yet)"

    exp_path = experiment_dir(abreka_dir, exp.id)
    findings_path = exp_path / "findings.md"
    findings_text = ""
    if findings_path.is_file():
        findings_text = findings_path.read_text().strip()
    if not findings_text:
        findings_text = "(no findings written yet)"

    return {
        "project_name": project_name,
        "lib_name": lib_name,
        "exp_id": exp.id,
        "status": exp.status.value,
        "hypothesis": exp.hypothesis or "(none)",
        "method": exp.method or "(none)",
        "parents_line": ", ".join(exp.parents) if exp.parents else "(none)",
        "tags_line": ", ".join(exp.tags) if exp.tags else "(none)",
        "created_at": exp.created_at.isoformat(),
        "completed_at": exp.completed_at.isoformat() if exp.completed_at else "(not completed)",
        "failure_reason": exp.failure_reason or "(n/a)",
        "metrics_table": _format_metrics_table(exp.metrics),
        "findings_text": findings_text,
        "file_listing": _list_experiment_files(exp_path),
        "primary_metric": config.experiment.primary_metric,
        "goal_text": goal_text,
        "parent_summaries": _format_parent_summaries(abreka_dir, exp.parents),
    }


def _render_system_prompt(context: dict[str, str]) -> str:
    template = _TEMPLATE_RESOURCE.read_text(encoding="utf-8")
    return render_template(template, context)


def _build_argv(
    *,
    system_prompt: str,
    chats_dir: Path,
    continue_: bool,
    resume: bool,
    messages: tuple[str, ...],
    exp_id: str,
) -> list[str]:
    argv = [
        "pi",
        "--append-system-prompt",
        system_prompt,
        "--session-dir",
        str(chats_dir),
    ]
    if continue_:
        argv.append("--continue")
    if resume:
        argv.append("--resume")
    if messages:
        argv.extend(messages)
    else:
        argv.append(
            f"Let's discuss experiment {exp_id}. I've given you its metadata, "
            "findings, and the project layout in the system prompt. What do you "
            "want to know first?"
        )
    return argv


@click.command("chat")
@click.argument("exp_id")
@click.option(
    "--continue", "continue_", is_flag=True,
    help="Continue the most recent chat for this experiment.",
)
@click.option(
    "--resume", is_flag=True,
    help="Pick a previous chat for this experiment to resume.",
)
@click.argument("messages", nargs=-1)
def chat_cmd(exp_id: str, continue_: bool, resume: bool, messages: tuple[str, ...]):
    """Open an interactive pi chat briefed on experiment EXP_ID.

    The chat runs from the project root so pi has access to shared library code
    and sibling experiments. Session files are persisted under
    abreka/experiments/<id>/chats/ so each experiment has its own history.
    """
    abreka_dir = find_abreka_dir()
    exp_id = normalize_exp_id(exp_id)
    exp = load_experiment(abreka_dir, exp_id)

    if not shutil.which("pi"):
        raise click.ClickException(
            "pi not found. Install it: https://github.com/badlogic/pi-mono"
        )

    context = _build_context(abreka_dir, exp)
    system_prompt = _render_system_prompt(context)

    chats_dir = experiment_dir(abreka_dir, exp_id) / "chats"
    chats_dir.mkdir(parents=True, exist_ok=True)

    argv = _build_argv(
        system_prompt=system_prompt,
        chats_dir=chats_dir,
        continue_=continue_,
        resume=resume,
        messages=messages,
        exp_id=exp_id,
    )

    project_root = abreka_dir.parent
    click.echo(f"Launching pi for experiment {exp_id} in {project_root}...\n")
    os.chdir(project_root)
    os.execvp("pi", argv)
