from __future__ import annotations

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

from amplitude_ai.testing import MockAmplitudeAI

PROVIDER_PACKAGES = (
    "openai",
    "anthropic",
    "google-genai",
    "mistralai",
    "boto3",
)


def _detect_providers_from_pyproject(path: Path) -> list[str]:
    """Parse pyproject.toml and check for provider packages in dependencies."""
    if not path.exists():
        return []

    try:
        import tomllib
    except ModuleNotFoundError:
        # Python < 3.11 fallback: simple string matching on raw text
        contents = path.read_text(encoding="utf-8")
        return [p for p in PROVIDER_PACKAGES if p in contents]

    with path.open("rb") as f:
        data = tomllib.load(f)

    dep_strings: list[str] = []
    project = data.get("project", {})
    dep_strings.extend(project.get("dependencies", []))
    for group_deps in project.get("optional-dependencies", {}).values():
        if isinstance(group_deps, list):
            dep_strings.extend(group_deps)

    dep_names = {dep.split(">")[0].split("<")[0].split("=")[0].split("[")[0].strip().lower() for dep in dep_strings}
    return [p for p in PROVIDER_PACKAGES if p.lower() in dep_names]


def run_doctor(cwd: Path, include_mock_check: bool = True) -> dict[str, Any]:
    checks: list[dict[str, Any]] = []

    has_api_key = bool(os.environ.get("AMPLITUDE_AI_API_KEY"))
    checks.append({
        "name": "env.AMPLITUDE_AI_API_KEY",
        "ok": has_api_key,
        "detail": "present" if has_api_key else "missing",
        **({"fix": "export AMPLITUDE_AI_API_KEY=your_project_api_key"} if not has_api_key else {}),
    })

    pyproject_path = cwd / "pyproject.toml"
    has_pyproject = pyproject_path.exists()
    checks.append({
        "name": "pyproject.toml",
        "ok": has_pyproject,
        "detail": "found" if has_pyproject else "missing",
        **({"fix": "Point your AI coding agent at the amplitude-ai.md guide shipped with the package"} if not has_pyproject else {}),
    })

    detected_providers = _detect_providers_from_pyproject(pyproject_path)
    checks.append({
        "name": "provider_dependency",
        "ok": len(detected_providers) > 0,
        "detail": (
            f"found {', '.join(detected_providers)}"
            if detected_providers
            else "no known provider package found"
        ),
        **({"fix": "pip install 'amplitude-ai[openai]'  # or [anthropic], [gemini], [bedrock], [mistral]"} if not detected_providers else {}),
    })

    if include_mock_check:
        try:
            mock = MockAmplitudeAI()
            mock.track_user_message(
                user_id="doctor-user",
                content="doctor probe",
                session_id="doctor-session",
            )
            user_events = mock.get_events("[Agent] User Message")
            checks.append({
                "name": "mock_event_delivery",
                "ok": len(user_events) > 0,
                "detail": (
                    f"captured {len(user_events)} event(s)"
                    if user_events
                    else "no events captured"
                ),
                **({"fix": "Ensure amplitude-analytics is installed: pip install amplitude-analytics"} if not user_events else {}),
            })

            flushed = mock.flush()
            checks.append({
                "name": "mock_flush_path",
                "ok": isinstance(flushed, list),
                "detail": (
                    f"flush returned {len(flushed)} item(s)"
                    if isinstance(flushed, list)
                    else "flush did not return a list"
                ),
                **({"fix": "Check AmplitudeAI initialization and ensure flush() is called before exit"} if not isinstance(flushed, list) else {}),
            })
        except Exception as error:
            checks.append(
                {
                    "name": "mock_event_delivery",
                    "ok": False,
                    "detail": f"mock check failed: {error}",
                }
            )

    ok = all(check["ok"] for check in checks)
    return {"ok": ok, "checks": checks}


def main() -> None:
    args = sys.argv[1:]
    include_mock_check = "--no-mock-check" not in args
    json_mode = "--json" in args
    result = run_doctor(Path.cwd(), include_mock_check=include_mock_check)

    if json_mode:
        print(json.dumps(result, indent=2))
    else:
        print("amplitude-ai doctor\n")
        for check in result["checks"]:
            mark = "\u2713" if check["ok"] else "\u2717"
            print(f"  {mark} {check['name']}: {check['detail']}")
            if check.get("fix"):
                print(f"    fix: {check['fix']}")
        print(f"\n{'All checks passed.' if result['ok'] else 'Some checks failed.'}")

    raise SystemExit(0 if result["ok"] else 1)
