#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import sys
import tempfile
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))

import server  # noqa: E402


def load_env_file(path: Path) -> None:
    if not path.exists():
        return
    for raw in path.read_text().splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip().strip('"').strip("'")
        if key and key not in os.environ:
            os.environ[key] = value


def key_prefix() -> str:
    key = os.environ.get("OMEMAI_API_KEY", "")
    if not key:
        return "missing"
    if key.startswith("omem_"):
        return key[:12] + "..."
    return "configured"


def quick_check() -> dict[str, Any]:
    os.environ.setdefault("OMEMAI_CLIENT_SOURCE", "gateway_doctor")
    cfg = server.Config.from_env()
    client = server.Client(cfg)
    projects = client.request("GET", "/agent/projects")
    health = client.request("GET", "/agent/health/me")
    return {
        "mode": "quick",
        "api_base": cfg.api_base,
        "key_prefix": cfg.key_prefix,
        "projects_count": len(projects or []),
        "agent_health": {"status": health.get("status"), "health_score": health.get("health_score"), "warnings": health.get("warnings", [])},
    }


def full_check(project_id: str | None) -> dict[str, Any]:
    os.environ.setdefault("OMEMAI_CLIENT_SOURCE", "gateway_doctor")
    cfg = server.Config.from_env()
    tmp_state = tempfile.NamedTemporaryFile(prefix="omemai-gateway-doctor-", suffix=".json", delete=True)
    os.environ["OMEMAI_GATEWAY_STATE"] = tmp_state.name
    gateway = server.Gateway()
    projects = gateway.projects({})["projects"]
    selected = project_id or cfg.default_project_id or (projects[0]["id"] if projects else None)
    if not selected:
        return {"mode": "full", "ok": False, "error": "No project_id available. Create a project or set OMEMAI_DEFAULT_PROJECT_ID."}
    started = gateway.start_task({"project_id": selected, "task": "OMemAI Gateway doctor full self-check"})
    user = gateway.user_message({"content": "OMemAI Gateway doctor test user message.", "metadata": {"source": "gateway_doctor"}})
    assistant = gateway.assistant_message({"content": "OMemAI Gateway doctor test assistant message.", "metadata": {"source": "gateway_doctor"}})
    checkpoint = gateway.checkpoint({"reason": "doctor", "wait": True})
    search = gateway.search({"query": "Gateway doctor", "scope": "current_project", "limit": 5})
    status = gateway.status({})
    return {
        "mode": "full",
        "ok": True,
        "project_id": selected,
        "session_id": started.get("session_id"),
        "user_message_ok": user.get("ok"),
        "assistant_message_ok": assistant.get("ok"),
        "checkpoint_status": checkpoint.get("status"),
        "search_results": len(search.get("results") or []),
        "agent_health": status.get("agent_health"),
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="OMemAI Gateway doctor")
    parser.add_argument("--env-file", default=str(ROOT / ".env"))
    parser.add_argument("--full", action="store_true", help="Run write-path validation: start session, save messages, checkpoint, search")
    parser.add_argument("--project-id", default=None)
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()
    load_env_file(Path(args.env_file))
    try:
        result = full_check(args.project_id) if args.full else quick_check()
        result["ok"] = bool(result.get("ok", True))
        result["key_prefix"] = key_prefix()
        if args.json:
            print(json.dumps(result, ensure_ascii=False, indent=2))
        else:
            print(f"gateway_doctor=ok mode={result['mode']} key_prefix={result['key_prefix']}")
            if "projects_count" in result:
                print(f"projects_count={result['projects_count']}")
            if result.get("agent_health"):
                h = result["agent_health"]
                print(f"agent_health_status={h.get('status')} score={h.get('health_score')}")
            if args.full:
                print(f"session_id={result.get('session_id')}")
                print(f"checkpoint_status={result.get('checkpoint_status')}")
                print(f"search_results={result.get('search_results')}")
        return 0 if result.get("ok") else 2
    except Exception as exc:
        msg = server.redact(str(exc))
        if args.json:
            print(json.dumps({"ok": False, "error": msg, "key_prefix": key_prefix()}, ensure_ascii=False, indent=2))
        else:
            print(f"gateway_doctor=failed key_prefix={key_prefix()} error={msg}")
        return 1


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