"""What a conversation's model holds: prompt granules, skills, tools, files touched."""

from __future__ import annotations

import re
from collections import Counter
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from collections.abc import Iterable, Mapping

    from beaver_gateway.agents.claude import ClaudeAgent
    from beaver_gateway.conversations.kinds import Kind

__all__ = ["bash_paths", "compose", "files_touched", "granules", "skill_sets"]

CHARS_PER_TOKEN = 3.2
READ_TOOLS = frozenset({"Read", "NotebookRead"})
WRITE_TOOLS = frozenset({"Write", "Edit", "MultiEdit", "NotebookEdit"})
SEARCH_TOOLS = frozenset({"Glob", "Grep"})
PATH_KEYS = ("file_path", "notebook_path", "path", "file", "filename")
_FRONTMATTER = re.compile(r"^---\s*\n(.*?)\n---", re.DOTALL)
# A shell command names files as quoted strings or bare words; only notes
# and plain text files count, code and binaries are the agent's own business.
_SHELL_TOKEN = re.compile(r'"([^"\n]+)"|\'([^\'\n]+)\'|([^\s"\'`;|&()<>]+)')
_NOTE_SUFFIXES = (".md", ".canvas", ".txt", ".csv", ".json", ".yaml", ".yml")
_SHELL_CD = re.compile(r'(?:^|&&|;|\|\|)\s*cd\s+("([^"]+)"|\'([^\']+)\'|(\S+))')
_SHELL_WRITES = re.compile(
    r"write_text|write_bytes|>>?\s*[\"\']?[^\s&|]|\btee\b|\bsed\s+-i|\bmv\b|\brm\b"
    r"|\btouch\b|\bcp\b|open\([^)]*[\"\'][wa]"
)


def bash_paths(command: str, cwd: Path | None) -> tuple[list[str], Path | None, bool]:
    """Files a shell command names, where it leaves the shell, whether it writes.

    Relative names resolve against ``cwd``, or against the directory of a
    leading ``cd`` in the same command. Absolute names stay as they are.
    """
    here = cwd
    for match in _SHELL_CD.finditer(command):
        target = match.group(2) or match.group(3) or match.group(4) or ""
        target = target.strip()
        if not target:
            continue
        here = Path(target) if target.startswith("/") else (here or Path()) / target
    found: list[str] = []
    for match in _SHELL_TOKEN.finditer(command):
        raw = (match.group(1) or match.group(2) or match.group(3) or "").strip()
        if not raw.lower().endswith(_NOTE_SUFFIXES) or "*" in raw:
            continue
        path = Path(raw) if raw.startswith("/") else (here or Path()) / raw
        text = str(path)
        if text not in found:
            found.append(text)
    return found, here, bool(_SHELL_WRITES.search(command))


def _tokens(chars: int) -> int:
    return round(chars / CHARS_PER_TOKEN)


def _relative(path: str, cwd: Path | None) -> str:
    if cwd is None:
        return path
    try:
        return str(Path(path).relative_to(cwd))
    except ValueError:
        return path


def _tool_uses(entries: Iterable[Mapping[str, Any]]) -> Iterable[tuple[str, Any, str]]:
    for entry in entries:
        message = entry.get("message")
        if entry.get("type") != "assistant" or not isinstance(message, dict):
            continue
        content = message.get("content")
        if not isinstance(content, list):
            continue
        stamp = str(entry.get("timestamp") or "")
        for block in content:
            if isinstance(block, dict) and block.get("type") == "tool_use":
                yield str(block.get("name") or "?"), block.get("input"), stamp


def files_touched(
    entries: Iterable[Mapping[str, Any]], *, cwd: Path | None = None
) -> tuple[list[dict[str, Any]], dict[str, int]]:
    """Files named in tool inputs, with how they were touched, plus tool counts."""
    files: dict[str, dict[str, Any]] = {}
    counts: Counter[str] = Counter()
    # The shell keeps its directory between calls; follow it.
    shell_cwd = cwd

    def touch(raw: str, how: str, stamp: str) -> None:
        key = _relative(raw, cwd)
        row = files.setdefault(
            key, {"path": key, "reads": 0, "writes": 0, "other": 0, "last_at": stamp}
        )
        row[how] += 1
        row["last_at"] = max(row["last_at"], stamp)

    for name, tool_input, stamp in _tool_uses(entries):
        counts[name] += 1
        if not isinstance(tool_input, dict):
            continue
        if name == "Bash" and isinstance(tool_input.get("command"), str):
            paths, shell_cwd, writes = bash_paths(tool_input["command"], shell_cwd)
            for raw in paths:
                touch(raw, "writes" if writes else "reads", stamp)
            continue
        raw = next(
            (v for k in PATH_KEYS if isinstance(v := tool_input.get(k), str) and v),
            None,
        )
        if raw is None or (name in SEARCH_TOOLS and "file_path" not in tool_input):
            continue
        if name in READ_TOOLS:
            touch(raw, "reads", stamp)
        elif name in WRITE_TOOLS:
            touch(raw, "writes", stamp)
        else:
            touch(raw, "other", stamp)
    ordered = sorted(files.values(), key=lambda r: r["last_at"], reverse=True)
    return ordered, dict(counts)


def granules(agent: ClaudeAgent, kind: Kind) -> dict[str, Any]:
    """The system prompt as its sources: tag, path, size, token estimate."""
    sources = agent.prompt_for(kind) if agent.serves(kind) else None
    rows: list[dict[str, Any]] = []
    if sources is None:
        chars = len(agent.system_prompt)
        if chars:
            rows.append(
                {
                    "tag": None,
                    "path": None,
                    "bytes": chars,
                    "tokens_est": _tokens(chars),
                }
            )
    else:
        for source in sources:
            tag, raw = source if isinstance(source, tuple) else (None, source)
            path = Path(raw)
            try:
                chars = len(path.read_text(encoding="utf-8"))
            except OSError:
                chars = 0
            rows.append(
                {
                    "tag": tag,
                    "path": _relative(str(path), agent.cwd),
                    "bytes": chars,
                    "tokens_est": _tokens(chars),
                }
            )
    total = sum(r["bytes"] for r in rows)
    return {"bytes": total, "tokens_est": _tokens(total), "granules": rows}


def _skill_meta(skill_md: Path) -> dict[str, str]:
    try:
        text = skill_md.read_text(encoding="utf-8")
    except OSError:
        return {}
    match = _FRONTMATTER.match(text)
    meta: dict[str, str] = {}
    if match:
        for line in match.group(1).splitlines():
            key, sep, value = line.partition(":")
            if sep:
                meta[key.strip()] = value.strip()
    return meta


def skill_sets(agent: ClaudeAgent, kind: Kind) -> list[dict[str, Any]]:
    """Skill-set directories the kind loads, with each skill's name and description."""
    out: list[dict[str, Any]] = []
    for directory in agent.skills_for(kind):
        skills = []
        try:
            children = sorted(p for p in directory.iterdir() if p.is_dir())
        except OSError:
            children = []
        for child in children:
            skill_md = child / "SKILL.md"
            if skill_md.is_file():
                meta = _skill_meta(skill_md)
                skills.append(
                    {
                        "name": meta.get("name") or child.name,
                        "description": meta.get("description", ""),
                        "path": _relative(str(child), agent.cwd),
                    }
                )
        out.append(
            {
                "set": directory.name,
                "path": _relative(str(directory), agent.cwd),
                "skills": skills,
            }
        )
    return out


def compose(
    agent: ClaudeAgent,
    kind: Kind,
    entries: Iterable[Mapping[str, Any]],
    *,
    context_tokens: int,
) -> dict[str, Any]:
    files, counts = files_touched(entries, cwd=agent.cwd)
    prompt = granules(agent, kind)
    skills = skill_sets(agent, kind)
    skills_chars = sum(
        len(s["name"]) + len(s["description"])
        for group in skills
        for s in group["skills"]
    )
    return {
        "agent": {
            "name": agent.name,
            "model": agent.model,
            "effort": agent.options.effort,
        },
        "kind": kind,
        "context_tokens": context_tokens,
        "prompt": prompt,
        "skills": skills,
        "skills_tokens_est": _tokens(skills_chars),
        "tools": {
            "allowed": list(agent.options.tools) if agent.options.tools else None,
            "disallowed": list(agent.options.disallowed_tools),
            "gateway": list(agent.gateway_tools),
            "mcps": [m.name for m in agent.expose_mcps],
        },
        "files": files,
        "tool_counts": counts,
        "history_tokens_est": max(
            0, context_tokens - prompt["tokens_est"] - _tokens(skills_chars)
        ),
    }