"""
Demo generator — creates sample .air.json records and compliance docs
so developers can experience AIR Blackbox without running Docker.

This is the zero-config quickstart: `air-blackbox demo`
"""

import json
import os
import uuid
from datetime import datetime, timedelta
from pathlib import Path


SAMPLE_RUNS = [
    {"model": "gpt-4o-mini", "provider": "openai", "tokens": {"prompt": 45, "completion": 182, "total": 227},
     "duration_ms": 1340, "status": "success", "tool_calls": []},
    {"model": "gpt-4o-mini", "provider": "openai", "tokens": {"prompt": 120, "completion": 350, "total": 470},
     "duration_ms": 2100, "status": "success", "tool_calls": ["web_search"]},
    {"model": "gpt-4o", "provider": "openai", "tokens": {"prompt": 200, "completion": 500, "total": 700},
     "duration_ms": 3200, "status": "success", "tool_calls": ["calculator", "db_query"]},
    {"model": "gpt-4o-mini", "provider": "openai", "tokens": {"prompt": 80, "completion": 0, "total": 80},
     "duration_ms": 5000, "status": "error", "tool_calls": []},
    {"model": "claude-3-5-sonnet", "provider": "anthropic", "tokens": {"prompt": 150, "completion": 400, "total": 550},
     "duration_ms": 1800, "status": "success", "tool_calls": ["code_interpreter"]},
    {"model": "gpt-4o-mini", "provider": "openai", "tokens": {"prompt": 60, "completion": 190, "total": 250},
     "duration_ms": 1100, "status": "success", "tool_calls": []},
    {"model": "gemini-1.5-pro", "provider": "google", "tokens": {"prompt": 300, "completion": 600, "total": 900},
     "duration_ms": 2800, "status": "success", "tool_calls": ["web_search", "calculator"]},
    {"model": "gpt-4o-mini", "provider": "openai", "tokens": {"prompt": 90, "completion": 220, "total": 310},
     "duration_ms": 1500, "status": "success", "tool_calls": []},
    {"model": "gpt-4o-mini", "provider": "openai", "tokens": {"prompt": 55, "completion": 0, "total": 55},
     "duration_ms": 30000, "status": "timeout", "tool_calls": []},
    {"model": "gpt-4o", "provider": "openai", "tokens": {"prompt": 180, "completion": 420, "total": 600},
     "duration_ms": 2600, "status": "success", "tool_calls": ["db_query"]},
]


def generate_demo_data(output_dir: str = ".") -> dict:
    """Generate sample .air.json records and compliance docs.

    Returns dict with summary of what was created.
    """
    runs_dir = os.path.join(output_dir, "runs")
    os.makedirs(runs_dir, exist_ok=True)

    created_files = []
    now = datetime.utcnow()

    # Generate .air.json records
    for i, sample in enumerate(SAMPLE_RUNS):
        run_id = str(uuid.uuid4())
        ts = (now - timedelta(hours=len(SAMPLE_RUNS) - i)).isoformat() + "Z"
        record = {
            "version": "1.0.0",
            "run_id": run_id,
            "trace_id": uuid.uuid4().hex[:16],
            "timestamp": ts,
            "model": sample["model"],
            "provider": sample["provider"],
            "request_vault_ref": f"vault://air-runs/{run_id}/request.json",
            "response_vault_ref": f"vault://air-runs/{run_id}/response.json",
            "request_checksum": f"sha256:{uuid.uuid4().hex[:12]}",
            "response_checksum": f"sha256:{uuid.uuid4().hex[:12]}",
            "tokens": sample["tokens"],
            "duration_ms": sample["duration_ms"],
            "status": sample["status"],
        }
        if sample["tool_calls"]:
            record["tool_calls"] = sample["tool_calls"]

        fpath = os.path.join(runs_dir, f"{run_id}.air.json")
        with open(fpath, "w") as f:
            json.dump(record, f, indent=2)
        created_files.append(fpath)

    # Generate RISK_ASSESSMENT.md template
    risk_path = os.path.join(output_dir, "RISK_ASSESSMENT.md")
    if not os.path.exists(risk_path):
        with open(risk_path, "w") as f:
            f.write(_RISK_TEMPLATE)
        created_files.append(risk_path)

    # Generate DATA_GOVERNANCE.md template
    dg_path = os.path.join(output_dir, "DATA_GOVERNANCE.md")
    if not os.path.exists(dg_path):
        with open(dg_path, "w") as f:
            f.write(_DATA_GOV_TEMPLATE)
        created_files.append(dg_path)

    models = list(set(s["model"] for s in SAMPLE_RUNS))
    providers = list(set(s["provider"] for s in SAMPLE_RUNS))
    total_tokens = sum(s["tokens"]["total"] for s in SAMPLE_RUNS)

    return {
        "runs_created": len(SAMPLE_RUNS),
        "files_created": len(created_files),
        "models": models,
        "providers": providers,
        "total_tokens": total_tokens,
        "runs_dir": runs_dir,
        "output_dir": output_dir,
    }


_RISK_TEMPLATE = """# Risk Assessment — [Your AI System Name]

> Auto-generated by AIR Blackbox. Fill in your system-specific details.
> This template covers EU AI Act Article 9 requirements.

## System Overview
- **System name**: [Your AI System]
- **Purpose**: [What the system does]
- **Risk classification**: [High-risk / Limited-risk / Minimal-risk]
- **Date**: [Date]
- **Author**: [Name]

## Identified Risks

| Risk | Likelihood | Impact | Mitigation | Status |
|------|-----------|--------|------------|--------|
| Model hallucination | Medium | High | Output validation, human review | Active |
| Prompt injection | Medium | High | AIR gateway injection detection | Active |
| PII leakage in prompts | Low | Critical | AIR prompt vault PII filtering | Active |
| Model version drift | Medium | Medium | AIR AI-BOM tracking | Planned |
| Runaway agent (cost) | Low | Medium | AIR guardrails cost limits | Active |
"""


_DATA_GOV_TEMPLATE = """# Data Governance — [Your AI System Name]

> Auto-generated by AIR Blackbox. Fill in your system-specific details.
> This template covers EU AI Act Article 10 requirements.

## Data Sources
| Source | Type | Consent | Retention | Notes |
|--------|------|---------|-----------|-------|
| User prompts | Runtime input | [Y/N] | [Period] | Stored in AIR vault |
| LLM responses | Runtime output | N/A | [Period] | Stored in AIR vault |
| RAG knowledge base | Reference data | [Y/N] | [Period] | [Source] |

## PII Handling
- **Detection**: AIR Blackbox OTel pipeline scans all prompts for PII
- **Redaction**: [Enabled/Disabled] via prompt vault
- **Storage**: Prompts stored in controlled vault (S3/MinIO), not third-party cloud

## Data Quality Measures
- Prompt injection detection via AIR gateway
- Response validation: [Describe your approach]
- RAG retrieval quality monitoring: [Describe your approach]

## Consent
- [Describe how user consent is obtained for data processing]
"""
