#!/usr/bin/env python3

import json
import os
import re
import shlex
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

from droid_transcript import collect_messages


def log_dir() -> Path:
    override = os.environ.get("KG_DROID_HOOK_LOG_DIR")
    if override:
        return Path(override).expanduser()
    return Path.home() / ".factory" / "kg-droid-hook"


def normalize_query(text: str) -> str:
    compact = re.sub(r"\s+", " ", text).strip()
    return compact[:200]


def phrase_queries(text: str) -> list[str]:
    phrases: list[str] = []
    for pattern in (r"`([^`]{2,80})`", r'"([^"\n]{2,80})"', r"'([^'\n]{2,80})'"):
        for match in re.findall(pattern, text):
            normalized = normalize_query(match)
            if normalized:
                phrases.append(normalized)
    return phrases


def build_queries(prompt: str, transcript_path: Path) -> list[str]:
    queries: list[str] = []

    prompt_query = normalize_query(prompt)
    if prompt_query:
        queries.append(prompt_query)
        queries.extend(phrase_queries(prompt_query))

    history = collect_messages(transcript_path)
    recent_user_messages = [item["text"] for item in history if item["role"] == "user"][
        -3:
    ]
    for message in recent_user_messages:
        normalized = normalize_query(message)
        if normalized:
            queries.append(normalized)
            queries.extend(phrase_queries(normalized))

    unique: list[str] = []
    seen: set[str] = set()
    for query in queries:
        lowered = query.lower()
        if lowered in seen:
            continue
        seen.add(lowered)
        unique.append(query)
        if len(unique) >= 6:
            break
    return unique


def find_kg_repo_root(start_path: Path) -> Path | None:
    current = start_path if start_path.is_dir() else start_path.parent
    for candidate in (current, *current.parents):
        cargo_toml = candidate / "Cargo.toml"
        if not cargo_toml.is_file():
            continue
        try:
            contents = cargo_toml.read_text(encoding="utf-8")
        except OSError:
            continue
        if 'name = "kg"' in contents:
            return candidate
    return None


def kg_command(project_cwd: str) -> tuple[list[str], Path | None]:
    raw = os.environ.get("KG_DROID_HOOK_KG_CMD", "").strip()
    if not raw:
        repo_root = find_kg_repo_root(Path(project_cwd).expanduser())
        if repo_root is not None:
            return ["cargo", "run", "--quiet", "--bin", "kg", "--"], repo_root
        return ["kg"], None

    parts = shlex.split(raw)
    return (parts or ["kg"]), None


def summarize_node(node: dict) -> str:
    properties = node.get("properties", {}) or {}
    description = properties.get("description") or ""
    key_facts = properties.get("key_facts") or []

    lines = [f"- {node.get('id')} | {node.get('name')} [{node.get('type')}]"]
    if description:
        lines.append(f"  Description: {description}")
    if key_facts:
        lines.append(f"  Facts: {'; '.join(key_facts[:2])}")
    return "\n".join(lines)


def lookup_graph_context(
    graph: str, queries: list[str], project_cwd: str
) -> tuple[str | None, str | None]:
    if not graph or not queries:
        return None, None

    limit = os.environ.get("KG_DROID_HOOK_FIND_LIMIT", "3")
    command_prefix, command_cwd = kg_command(project_cwd)
    command = command_prefix + [
        "graph",
        graph,
        "node",
        "find",
        *queries,
        "--json",
        "--limit",
        limit,
    ]

    try:
        result = subprocess.run(
            command,
            capture_output=True,
            text=True,
            check=True,
            timeout=8,
            cwd=command_cwd,
        )
    except (
        subprocess.CalledProcessError,
        FileNotFoundError,
        subprocess.TimeoutExpired,
    ) as exc:
        return None, f"KG context lookup failed for graph `{graph}`: {exc}"

    try:
        payload = json.loads(result.stdout)
    except json.JSONDecodeError as exc:
        return (
            None,
            f"KG context lookup returned invalid JSON for graph `{graph}`: {exc}",
        )

    nodes: list[dict] = []
    seen: set[str] = set()
    for query_result in payload.get("queries", []):
        for node in query_result.get("nodes", []):
            node_id = node.get("id")
            if not node_id or node_id in seen:
                continue
            seen.add(node_id)
            nodes.append(node)
            if len(nodes) >= 5:
                break
        if len(nodes) >= 5:
            break

    if not nodes:
        return None, None

    summary_lines = [f"Relevant kg graph context from `{graph}`:"]
    summary_lines.extend(summarize_node(node) for node in nodes)
    return "\n".join(summary_lines), None


def main() -> int:
    try:
        payload = json.load(sys.stdin)
    except json.JSONDecodeError as exc:
        print(f"invalid hook input: {exc}", file=sys.stderr)
        return 1

    out_dir = log_dir()
    out_dir.mkdir(parents=True, exist_ok=True)

    prompt = payload.get("prompt", "")
    transcript_path = Path(payload.get("transcript_path", "")).expanduser()
    project_cwd = str(payload.get("cwd") or os.getcwd())
    graph = os.environ.get("KG_DROID_HOOK_GRAPH", "").strip()
    queries = build_queries(prompt, transcript_path)
    additional_context, context_error = lookup_graph_context(
        graph, queries, project_cwd
    )

    entry = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "hook_event_name": payload.get("hook_event_name"),
        "session_id": payload.get("session_id"),
        "cwd": payload.get("cwd"),
        "transcript_path": payload.get("transcript_path"),
        "prompt": prompt,
        "graph": graph or None,
        "queries": queries,
        "context_added": additional_context is not None,
        "additional_context": additional_context,
        "context_error": context_error,
    }

    with (out_dir / "prompts.jsonl").open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(entry, ensure_ascii=True) + "\n")

    if additional_context:
        print(
            json.dumps(
                {
                    "hookSpecificOutput": {
                        "hookEventName": "UserPromptSubmit",
                        "additionalContext": additional_context,
                    }
                },
                ensure_ascii=True,
            )
        )

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
