#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# Copyright 2024-2026 Lukasz Krzemien (biuro@softspark.eu)
# Source: https://github.com/softspark/ai-toolkit

"""Generate native Codex custom-agent TOML files."""

from __future__ import annotations

import json
import os
import sys
import tempfile
import tomllib
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from emission import agents_dir
from frontmatter import frontmatter_field

AGENT_PREFIX = "ai-toolkit-"
MANAGED_MARKER = "# ai-toolkit-managed: codex-agent"
REQUIRED_FIELDS = {"name", "description", "developer_instructions"}


def _agent_body(agent_file: Path) -> str:
    text = agent_file.read_text(encoding="utf-8")
    lines = text.splitlines(keepends=True)
    body = text
    if lines and lines[0].rstrip("\r\n") == "---":
        for index, line in enumerate(lines[1:], start=1):
            if line.rstrip("\r\n") == "---":
                body = "".join(lines[index + 1:])
                break
    return body.lstrip("\r\n").rstrip() + "\n"


def _toml_string(value: str) -> str:
    return json.dumps(value, ensure_ascii=False)


def _render_agent(agent_file: Path) -> str:
    name = frontmatter_field(agent_file, "name")
    description = frontmatter_field(agent_file, "description")
    body = _agent_body(agent_file)
    return "\n".join(
        [
            "# Generated by ai-toolkit. Do not edit.",
            MANAGED_MARKER,
            f"name = {_toml_string(name)}",
            f"description = {_toml_string(description)}",
            f"developer_instructions = {_toml_string(body)}",
            "",
        ]
    )


def _is_managed(path: Path) -> bool:
    if path.is_symlink() or not path.is_file():
        return False
    try:
        lines = path.read_text(encoding="utf-8").splitlines()
    except (OSError, UnicodeError):
        return False
    return MANAGED_MARKER in lines[:3]


def _warn_preserved(path: Path, reason: str) -> None:
    print(
        f"Warning: preserving user Codex agent '{path}': {reason}",
        file=sys.stderr,
    )


def _unmanaged_agent_names(output_dir: Path) -> set[str]:
    names: set[str] = set()
    for path in sorted(output_dir.glob("*.toml")):
        if path.is_symlink():
            _warn_preserved(path, "path is a symlink")
            continue
        if _is_managed(path):
            continue
        try:
            data = tomllib.loads(path.read_text(encoding="utf-8"))
        except (OSError, UnicodeError, tomllib.TOMLDecodeError) as error:
            _warn_preserved(path, f"cannot parse TOML ({error})")
            continue
        name = data.get("name")
        if isinstance(name, str) and name.strip():
            names.add(name)
    return names


def _cleanup_stale(output_dir: Path, expected: set[str]) -> int:
    removed = 0
    for output in sorted(output_dir.glob(f"{AGENT_PREFIX}*.toml")):
        if output.is_symlink():
            continue
        if output.name in expected or not _is_managed(output):
            continue
        output.unlink()
        removed += 1
    return removed


def _assert_safe_output_paths(base: Path, output_dir: Path) -> None:
    if base.is_symlink():
        raise RuntimeError(f"Refusing symlinked Codex config directory: {base}")
    if output_dir.is_symlink():
        raise RuntimeError(f"Refusing symlinked Codex agents directory: {output_dir}")


def _prepare_output_dir(base: Path) -> Path:
    output_dir = base / "agents"
    _assert_safe_output_paths(base, output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    _assert_safe_output_paths(base, output_dir)
    return output_dir


def _validate_agent_toml(rendered: str, source: Path) -> None:
    data = tomllib.loads(rendered)
    if set(data) != REQUIRED_FIELDS:
        raise ValueError(f"Invalid Codex agent fields rendered from {source}")
    if any(not isinstance(data[field], str) or not data[field] for field in REQUIRED_FIELDS):
        raise ValueError(f"Empty Codex agent field rendered from {source}")


def _build_write_plan(
    output_dir: Path,
    unmanaged_names: set[str],
) -> tuple[list[tuple[Path, str]], set[str]]:
    plan: list[tuple[Path, str]] = []
    expected: set[str] = set()

    for agent_file in sorted(agents_dir.glob("*.md")):
        name = frontmatter_field(agent_file, "name")
        description = frontmatter_field(agent_file, "description")
        if not name or not description:
            continue
        if name in unmanaged_names:
            continue
        filename = f"{AGENT_PREFIX}{name}.toml"
        expected.add(filename)
        output = output_dir / filename
        if output.is_symlink():
            _warn_preserved(output, "destination is a symlink")
            continue
        if output.exists() and not _is_managed(output):
            continue
        rendered = _render_agent(agent_file)
        _validate_agent_toml(rendered, agent_file)
        plan.append((output, rendered))
    return plan, expected


def _discard_temp(path: Path, output_dir: Path) -> None:
    if path.parent != output_dir:
        return
    try:
        path.unlink(missing_ok=True)
    except OSError:
        pass


def _stage_agent(output_dir: Path, output: Path, rendered: str) -> Path:
    fd, temp_name = tempfile.mkstemp(
        dir=output_dir,
        prefix=f".{output.name}.",
        suffix=".tmp",
    )
    temp_path = Path(temp_name)
    try:
        handle = os.fdopen(fd, "w", encoding="utf-8")
        fd = -1
        with handle:
            handle.write(rendered)
            handle.flush()
            os.fsync(handle.fileno())
        _validate_agent_toml(temp_path.read_text(encoding="utf-8"), output)
        return temp_path
    except Exception:
        if fd >= 0:
            os.close(fd)
        _discard_temp(temp_path, output_dir)
        raise


def _stage_all(
    output_dir: Path,
    plan: list[tuple[Path, str]],
) -> list[tuple[Path, Path]]:
    staged: list[tuple[Path, Path]] = []
    try:
        for output, rendered in plan:
            staged.append((_stage_agent(output_dir, output, rendered), output))
        return staged
    except Exception:
        for temp_path, _ in staged:
            _discard_temp(temp_path, output_dir)
        raise


def _replace_all(
    base: Path,
    output_dir: Path,
    staged: list[tuple[Path, Path]],
) -> int:
    written = 0
    _assert_safe_output_paths(base, output_dir)
    for temp_path, output in staged:
        if output.is_symlink():
            _warn_preserved(output, "destination became a symlink during generation")
            continue
        if output.exists() and not _is_managed(output):
            _warn_preserved(output, "destination became user-owned during generation")
            continue
        os.replace(temp_path, output)
        written += 1
    return written


def generate(target_dir: Path, config_root: Path | None = None) -> tuple[int, int]:
    base = config_root if config_root is not None else target_dir / ".codex"
    output_dir = _prepare_output_dir(base)
    unmanaged_names = _unmanaged_agent_names(output_dir)
    plan, expected = _build_write_plan(output_dir, unmanaged_names)
    staged = _stage_all(output_dir, plan)
    try:
        written = _replace_all(base, output_dir, staged)
        _assert_safe_output_paths(base, output_dir)
    finally:
        for temp_path, _ in staged:
            _discard_temp(temp_path, output_dir)
    removed = _cleanup_stale(output_dir, expected)
    return written, removed


def main() -> None:
    target = Path(sys.argv[1]) if len(sys.argv) > 1 else Path.cwd()
    written, removed = generate(target)
    print(f"Generated: .codex/agents/ ({written} agents, {removed} stale removed)")


if __name__ == "__main__":
    main()