from __future__ import annotations

import argparse
import json
from pathlib import Path

from amplitude_ai.mcp.contract import GENERATED_FILES, MCP_PROMPTS, MCP_RESOURCES, MCP_TOOLS

ROOT = Path(__file__).resolve().parent.parent


def _read_version() -> str:
    pyproject = (ROOT / "pyproject.toml").read_text(encoding="utf-8")
    for line in pyproject.splitlines():
        if line.startswith("version = "):
            return line.split("=", maxsplit=1)[1].strip().strip('"')
    return "unknown"


def _read_event_names() -> list[str]:
    catalog = json.loads((ROOT / "amplitude_ai" / "data" / "agent_event_catalog.json").read_text(encoding="utf-8"))
    events = catalog.get("events", [])
    if not isinstance(events, list):
        return []
    names = [event.get("event_type") for event in events if isinstance(event, dict)]
    return sorted([name for name in names if isinstance(name, str)])


def _render_agents_md(version: str, event_names: list[str]) -> str:
    generated_header = "<!-- GENERATED FILE: do not edit manually -->"
    tools = ", ".join([f"`{name}`" for name in MCP_TOOLS.values()])
    resources = ", ".join([f"`{name}`" for name in MCP_RESOURCES.values()])
    event_lines = "\n".join([f"- `{name}`" for name in event_names])
    return f"""{generated_header}
# AGENTS.md

Package: `amplitude-ai` v{version}

## Install
- `pip install amplitude-ai`
- Optional MCP support: `pip install 'amplitude-ai[mcp]'`

## Decision Tree
- **Default: use `ai.agent(...).session(...)` with provider wrappers** — gives you every event type, per-user analytics, session enrichment, quality scoring.
- Already have provider client objects: use `wrap()` to instrument them.
- Multiple agents collaborating: use `session.run_as(child_agent)` for automatic identity propagation.
- Need explicit tool telemetry: use `@tool`.
- Cannot modify call sites at all: use `patch()` for aggregate-only monitoring (no per-user analytics).
- Need guided agent integration: use MCP prompt `instrument_app`.

## Canonical Patterns
- zero-code patching
- wrap-openai
- bound-agent-session
- multi-agent-run-as
- tool-decorator
- fastapi-middleware

## MCP Surface
- Tools: {tools}
- Resources: {resources}
- Prompt: `{MCP_PROMPTS["instrument_app"]}`

## Gotchas
- Keep `AMPLITUDE_AI_API_KEY` set in runtime env.
- Use `MockAmplitudeAI` for deterministic tests.
- Use `unpatch()` in tests to prevent cross-test leakage.

## Testing
- `uv run pytest langley/amplitude_ai/tests`
- `uv run python scripts/generate_agent_docs.py --check`

## CLI
- `amplitude-ai` — print instrumentation prompt for AI coding agents; subcommands: mcp, doctor, status, --print-guide
- `amplitude-ai-instrument` — zero-code auto-patching wrapper (run your app with automatic instrumentation)
- `amplitude-ai-doctor` — validate environment, deps, and event pipeline
- `amplitude-ai-status` — show SDK version, installed providers, and env config
- `amplitude-ai-mcp` — start the MCP server for AI coding agents
- `amplitude-ai-register-catalog` — register event schema in your Amplitude Data Catalog

## Examples
- `examples/zero_code_example.py`
- `examples/wrap_openai_example.py`
- `examples/multi_agent_example.py`
- `examples/framework_integration_example.py`
- `examples/real_openai_example.py` (requires OPENAI_API_KEY)

## Extended Reference
- `llms-full.txt` — Full API signatures and canonical patterns for LLM agents

## Instrumentation Guide
- `amplitude-ai.md` — self-contained 4-phase instrumentation guide for any AI coding agent

## Event Schema (names)
{event_lines}
"""


def _render_llms_txt(version: str, event_names: list[str]) -> str:
    generated_header = "<!-- GENERATED FILE: do not edit manually -->"
    tools = "\n".join(MCP_TOOLS.values())
    resources = "\n".join(MCP_RESOURCES.values())
    events = "\n".join(event_names)
    return f"""{generated_header}
# llms.txt
package=amplitude-ai
version={version}

[mcp.tools]
{tools}

[mcp.resources]
{resources}

[events]
{events}
"""

def _render_llms_full_txt(version: str) -> str:
    generated_header = "# llms-full.txt"
    return f"""{generated_header}
# amplitude-ai {version} — Detailed API Reference for LLM Agents
# Use this file for instrumentation guidance. See llms.txt for discovery.

## Core API

### AmplitudeAI(amplitude)
Initialize the SDK. Required entry point.
```
from amplitude import Amplitude
from amplitude_ai import AmplitudeAI
amplitude = Amplitude("api-key")
ai = AmplitudeAI(amplitude=amplitude, content_mode="full")
```

### patch(amplitude)
Zero-code instrumentation. Monkey-patches all detected provider SDKs.
```
from amplitude_ai import patch, unpatch
patch(amplitude=amplitude)
unpatch()
```

### wrap(client, amplitude)
Convert an existing provider client into an instrumented wrapper. Supports both sync and async clients.
```
from amplitude_ai import wrap
instrumented = wrap(openai_client, amplitude=amplitude)  # works with openai.OpenAI and openai.AsyncOpenAI
```

### ai.agent(agent_id, user_id)
Create a bound agent for user/session lineage.
```
agent = ai.agent("my-agent", user_id="u1")
child = agent.child("sub-agent")
with agent.session(session_id="s1") as session:
    ...  # LLM calls tracked automatically
    with session.run_as(child) as cs:
        ...  # delegate to child agent
```

### @tool / @observe
Decorators for automatic function tracking.
```
from amplitude_ai import tool, observe

@tool
def search_products(query: str) -> dict:
    return {{'items': []}}

@observe
async def handle_request(request):
    ...  # Emits [Agent] Span events with session lifecycle
```

### MockAmplitudeAI
Deterministic testing helper. Drop-in replacement.
```
from amplitude_ai import MockAmplitudeAI
mock = MockAmplitudeAI()
mock.get_events("[Agent] AI Response")  # → list of events
mock.flush()  # → list of events
```

## Provider Wrappers

| Provider     | Import                                   | Streaming | Tool Calls | TTFB | Cache Tokens |
|--------------|------------------------------------------|-----------|------------|------|--------------|
| OpenAI       | `from amplitude_ai import OpenAI`        | Yes       | Yes        | Yes  | Yes          |
| AsyncOpenAI  | `from amplitude_ai import AsyncOpenAI`   | Yes       | Yes        | Yes  | Yes          |
| Anthropic    | `from amplitude_ai import Anthropic`     | Yes       | Yes        | Yes  | Yes          |
| Gemini       | `from amplitude_ai import Gemini`        | Yes       | No         | No   | No           |
| AzureOpenAI  | `from amplitude_ai import AzureOpenAI`   | Yes       | Yes        | Yes  | No           |
| Bedrock      | `from amplitude_ai import Bedrock`       | Yes       | Yes        | No   | No           |
| Mistral      | `from amplitude_ai import Mistral`       | Yes       | No         | No   | No           |

## Canonical Patterns

### 1. Zero-code (patch)
```python
from amplitude import Amplitude
from amplitude_ai import AmplitudeAI, patch
amplitude = Amplitude("api-key")
ai = AmplitudeAI(amplitude=amplitude)
patch(amplitude=amplitude)
```

### 2. Wrap existing client
```python
from amplitude_ai import wrap
instrumented = wrap(existing_client, amplitude=amplitude)
```

### 3. Provider wrapper
```python
from amplitude_ai import AmplitudeAI, OpenAI
client = OpenAI(amplitude=ai, api_key="sk-...")
response = client.chat.completions.create(model="gpt-4o", messages=[...], amplitude_user_id="u1")
```

### 4. Bound agent + session
```python
agent = ai.agent("assistant", user_id="u1")
with agent.session(session_id="s1") as session:
    response = client.chat.completions.create(...)
```

### 5. FastAPI middleware
```python
from amplitude_ai.middleware import AmplitudeAIMiddleware
app.add_middleware(AmplitudeAIMiddleware, ai=ai)
```

### 6. Multi-agent orchestration with session.run_as()
```python
orchestrator = ai.agent("orchestrator", user_id="u1")
researcher = orchestrator.child("researcher")
with orchestrator.session(session_id="s1") as s:
    # Provider calls inside run_as are automatically tagged with the child's agent_id
    with s.run_as(researcher) as cs:
        response = client.chat.completions.create(model="gpt-4o", messages=[...])
# run_as shares session_id, trace_id, turn counter; does NOT emit Session End
```

## MCP Tools

- `get_event_schema(event_type?)` — Return event schema and property definitions
- `get_integration_pattern(pattern_id?)` — Return canonical instrumentation patterns
- `validate_setup()` — Check required environment variables
- `suggest_instrumentation(framework?, provider?, content_tier?)` — Value-first setup guidance with content-tier and privacy defaults
- `validate_file(source, language?)` — Detect uninstrumented LLM call sites
- `search_docs(query, max_results?)` — Search README and API reference by keyword

## CLI

- `amplitude-ai` — Print setup prompt for AI coding agents; `--print-guide` outputs full guide
- `amplitude-ai mcp` — Start the MCP server for AI coding agents
- `amplitude-ai doctor [--json] [--no-mock-check]` — Validate environment and event pipeline
- `amplitude-ai status [--json]` — Show SDK version, installed providers, and env config
- `amplitude-ai-register-catalog` — Push event catalog to Amplitude Data

## Common Errors

- "No events captured" → Ensure session context wraps your LLM calls
- "patch() drops events silently" → patch() requires active SessionContext
- "flush() timeout" → Call ai.flush() before process exit in serverless
- Import errors → Install optional deps: pip install 'amplitude-ai[openai]'
"""


def _render_mcp_schema(version: str) -> str:
    payload = {
        "generated": True,
        "package": "amplitude-ai",
        "version": version,
        "prompt": MCP_PROMPTS["instrument_app"],
        "tools": list(MCP_TOOLS.values()),
        "resources": list(MCP_RESOURCES.values()),
    }
    return f"{json.dumps(payload, indent=2)}\n"


def _write_or_check(path: Path, content: str, check: bool) -> bool:
    existing = path.read_text(encoding="utf-8") if path.exists() else ""
    if existing == content:
        return False
    if not check:
        path.write_text(content, encoding="utf-8")
    return True


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate AGENTS.md and llms.txt for amplitude-ai")
    parser.add_argument("--check", action="store_true", help="Fail if generated files are stale")
    args = parser.parse_args()

    version = _read_version()
    event_names = _read_event_names()
    agents_md = _render_agents_md(version, event_names)
    llms_txt = _render_llms_txt(version, event_names)
    llms_full_txt = _render_llms_full_txt(version)
    mcp_schema = _render_mcp_schema(version)

    stale = False
    stale = _write_or_check(ROOT / GENERATED_FILES["agents"], agents_md, args.check) or stale
    stale = _write_or_check(ROOT / GENERATED_FILES["llms"], llms_txt, args.check) or stale
    stale = _write_or_check(ROOT / GENERATED_FILES["llms_full"], llms_full_txt, args.check) or stale
    stale = _write_or_check(ROOT / GENERATED_FILES["mcp_schema"], mcp_schema, args.check) or stale

    if args.check and stale:
        raise SystemExit("Generated docs are stale. Run: uv run python scripts/generate_agent_docs.py")


if __name__ == "__main__":
    main()
