"""
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)

    # Generate sample agent Python file for code scanner
    agent_path = os.path.join(output_dir, "sample_agent.py")
    if not os.path.exists(agent_path):
        with open(agent_path, "w") as f:
            f.write(_SAMPLE_AGENT)
        created_files.append(agent_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]
"""


_SAMPLE_AGENT = '''"""
Sample AI Agent — generated by AIR Blackbox demo.

This agent demonstrates common patterns that the compliance scanner checks.
Some are good practices (logging, error handling). Others are gaps
(no user identity binding, no action boundaries).

Run: air-blackbox comply -v  to see the full compliance report.
"""
import logging
from dataclasses import dataclass
from typing import Optional
from openai import OpenAI

logger = logging.getLogger(__name__)


@dataclass
class AgentConfig:
    """Configuration for the sample agent."""
    model: str = "gpt-4o-mini"
    max_tokens: int = 1000
    max_retries: int = 3


def run_agent(query: str, config: Optional[AgentConfig] = None) -> str:
    """Run the AI agent with a user query.

    Args:
        query: The user's input query
        config: Optional agent configuration

    Returns:
        The agent's response text
    """
    config = config or AgentConfig()
    client = OpenAI()

    # ✅ GOOD: Error handling around LLM calls
    try:
        logger.info(f"Agent processing query: {query[:50]}...")
        response = client.chat.completions.create(
            model=config.model,
            max_tokens=config.max_tokens,
            messages=[{"role": "user", "content": query}],
        )
        result = response.choices[0].message.content
        logger.info(f"Agent completed. Tokens: {response.usage.total_tokens}")
        return result

    except Exception as e:
        # ✅ GOOD: Fallback on error
        logger.error(f"Agent failed: {e}")
        return f"Sorry, I encountered an error: {e}"


def send_email(to: str, subject: str, body: str) -> bool:
    """Send an email on behalf of the user.

    ⚠️ GAP: No user_id tracking — who authorized this?
    ⚠️ GAP: No action audit trail — was this logged?
    ⚠️ GAP: No action boundaries — can the agent send to anyone?
    """
    # This action is NOT logged, NOT tied to a user, NOT bounded
    print(f"Sending email to {to}: {subject}")
    return True


def search_web(query: str) -> str:
    """Search the web for information."""
    logger.info(f"Web search: {query}")
    return f"Results for: {query}"


if __name__ == "__main__":
    result = run_agent("What is the EU AI Act?")
    print(result)
'''
