from __future__ import annotations

from datetime import datetime, timezone
import json
import os
import re
import tempfile
from typing import Any

from puppet_theater.models import TheaterSession

"""
Trace logging and export system for the AI Puppet Theater.

This module is responsible for collecting, sanitizing, normalizing, and exporting
runtime execution traces of a theater session.

It captures all important system events such as:
- Actor responses and intents
- Director decisions
- Tool usage (requests and results)
- Audience interactions
- Backend and validation metadata
- Errors and fallback behavior

Key responsibilities:
1. Add structured trace events during a TheaterSession.
2. Sanitize sensitive information (secrets, tokens, file paths).
3. Normalize heterogeneous event formats into a consistent schema.
4. Generate human-readable summaries for debugging and inspection.
5. Export full session traces as JSON (in-memory or file-based).

This module is primarily used for debugging, observability, and replay
of AI-driven theater sessions.
"""


APP_NAME = "AI Puppet Theater"
TRACE_VERSION = "1.0"
_MAX_TEXT_CHARS = 500
_STANDARD_EVENT_KEYS = (
    "event_type",
    "timestamp",
    "step",
    "beat_index",
    "story_phase",
    "speaker",
    "director_decision",
    "actor_intent",
    "actor_memory_update",
    "tool_request",
    "tool_result",
    "audience_action",
    "backend_metadata",
    "latency_ms",
    "validation_status",
    "fallback_used",
    "fallback_reason",
    "summary",
)
_PRIVATE_PATH_PATTERNS = [
    re.compile(r"/Users/[^\s\"'<>:]+(?:/[^\s\"'<>:]+)*"),
    re.compile(r"/private/(?:tmp|var)/[^\s\"'<>:]+(?:/[^\s\"'<>:]+)*"),
    re.compile(r"/tmp/[^\s\"'<>:]+(?:/[^\s\"'<>:]+)*"),
]


def add_trace_event(session: TheaterSession, event_type: str, **fields: Any) -> None:
    """
    Adds a structured event to the session trace log.

    Automatically attaches metadata such as timestamp and step number,
    and sanitizes all values before storing them in the session.

    Used to record runtime events like actor responses, director decisions,
    tool calls, and system actions.
    """
    event: dict[str, Any] = {
        "event_type": sanitize_value(event_type),
        "type": sanitize_value(event_type),
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "step": len(session.trace_events) + 1,
    }
    if session.beat_index is not None:
        event["beat_index"] = session.beat_index
    for key, value in fields.items():
        if value is not None:
            event[key] = sanitize_value(value)
    session.trace_events.append(event)


def export_trace(session: TheaterSession | None) -> dict[str, Any] | None:
    """
    Builds a complete structured snapshot of the current theater session.

    Includes session metadata, backend configuration, and all normalized
    trace events in a consistent export format suitable for logging or analysis.
    """
    if session is None:
        return None

    actor_model_id = session.backend_model_id
    director_model_id = _model_id_for_mode(session.director_mode)
    active_model_id = actor_model_id or director_model_id
    return {
        "app_name": APP_NAME,
        "trace_version": TRACE_VERSION,
        "session_id": sanitize_value(session.session_id),
        "created_at": sanitize_value(session.created_at),
        "premise": sanitize_value(session.premise),
        "title": sanitize_value(session.show_title),
        "setting": sanitize_value(session.setting),
        "show_length_mode": sanitize_value(session.show_length_mode),
        "min_beats": session.min_beats,
        "target_beats": session.target_beats,
        "max_beats": session.max_beats,
        "active_backend": sanitize_value(session.backend_name),
        "actor_backend": sanitize_value(session.backend_name),
        "director_mode": sanitize_value(session.director_mode),
        "director_backend": sanitize_value(session.director_mode),
        "model_id": sanitize_value(active_model_id),
        "actor_model_id": sanitize_value(actor_model_id),
        "director_model_id": sanitize_value(director_model_id),
        "events": normalize_trace_events(session),
    }


def render_trace_json(session: TheaterSession | None) -> str:
    """
    Converts the exported trace into a pretty-printed JSON string.

    Useful for debugging, logs, or displaying the full session state
    in a human-readable format.
    """
    payload = export_trace(session)
    if payload is None:
        return "No trace events yet."
    return json.dumps(payload, indent=2, sort_keys=False)


def write_trace_json_file(session: TheaterSession | None) -> str | None:
    """
    Writes the full trace export to a temporary JSON file.

    The file is stored in the system temp directory and named using the
    session ID. Used for downloading or offline inspection of traces.
    """
    payload = export_trace(session)
    if payload is None:
        return None

    safe_session_id = re.sub(r"[^a-zA-Z0-9_-]+", "-", str(payload["session_id"]))[:48] or "session"
    filename = f"ai-puppet-theater-trace-{safe_session_id}.json"
    path = os.path.join(tempfile.gettempdir(), filename)
    with open(path, "w", encoding="utf-8") as trace_file:
        json.dump(payload, trace_file, indent=2, ensure_ascii=False)
        trace_file.write("\n")
    return path


def normalize_trace_events(session: TheaterSession) -> list[dict[str, Any]]:
    """
    Converts raw trace events into a standardized format.

    Handles both structured dict events and legacy string events,
    ensuring all events conform to a consistent schema for downstream use.
    """
    normalized: list[dict[str, Any]] = []
    for index, raw_event in enumerate(session.trace_events, start=1):
        if isinstance(raw_event, dict):
            normalized.append(_normalize_event_dict(raw_event, index))
            continue

        normalized.append(
            {
                "event_type": "legacy_event",
                "step": index,
                "message": sanitize_value(raw_event),
                "summary": sanitize_value(raw_event),
            }
        )
    return normalized


def render_trace_summary(session: TheaterSession | None, max_events: int = 14) -> str:
    """
    Generates a human-readable summary of recent trace events.

    Useful for quick debugging or UI display. Shows only the most recent
    events with concise one-line explanations.
    """
    if session is None:
        return "No trace events yet."

    events = normalize_trace_events(session)
    if not events:
        return "No trace events yet."

    lines = [
        f"{session.show_title} ({session.session_id})",
        (
            f"{session.show_length_mode} show: {session.min_beats}/"
            f"{session.target_beats}/{session.max_beats} beats, "
            f"actor backend={session.backend_name}, director={session.director_mode}"
        ),
        "",
    ]
    for event in events[-max_events:]:
        step = event.get("step", "?")
        event_type = event.get("event_type", "event")
        beat = event.get("beat_index")
        prefix = f"{step}. {event_type}"
        if beat is not None:
            prefix += f" [beat {beat}]"
        lines.append(f"{prefix}: {_event_summary(event)}")
    return "\n".join(lines)


def _normalize_event_dict(raw_event: dict[str, Any], index: int) -> dict[str, Any]:
    """
    Normalizes a raw event dictionary into a structured trace format.

    Extracts known fields, groups related metadata (director decisions,
    tool usage, actor updates), and ensures output conforms to the
    standard event schema.
    """
    cleaned = {str(key): sanitize_value(value) for key, value in raw_event.items() if value is not None}
    event_type = str(cleaned.pop("event_type", cleaned.pop("type", "event")) or "event")
    normalized: dict[str, Any] = {
        "event_type": sanitize_value(event_type),
        "step": cleaned.pop("step", index),
    }
    if "timestamp" in cleaned:
        normalized["timestamp"] = cleaned.pop("timestamp")

    for key in (
        "beat_index",
        "story_phase",
        "speaker",
        "latency_ms",
        "validation_status",
        "fallback_used",
        "fallback_reason",
    ):
        if key in cleaned:
            normalized[key] = cleaned.pop(key)

    if event_type == "director_decision":
        normalized["director_decision"] = {
            key: cleaned.pop(key)
            for key in (
                "beat_type",
                "reason_summary",
                "progress",
                "uses_prop",
                "reveal_secret",
                "should_end_scene",
                "min_beats",
                "target_beats",
                "max_beats",
            )
            if key in cleaned
        }
    if event_type == "actor_intent" and "intent" in cleaned:
        normalized["actor_intent"] = cleaned.pop("intent")
    if event_type == "actor_memory_update" and "memory_update" in cleaned:
        normalized["actor_memory_update"] = cleaned.pop("memory_update")
    if event_type in {"tool_requested", "tool_executed", "tool_ignored"}:
        normalized["tool_request"] = {
            key: cleaned.pop(key)
            for key in ("tool_name", "reason", "arguments")
            if key in cleaned
        }
    if event_type == "tool_result":
        normalized["tool_result"] = {
            key: cleaned.pop(key)
            for key in ("tool_name", "result", "stage_effect")
            if key in cleaned
        }
    if event_type == "audience_action":
        action = cleaned.pop("audience_action", None)
        normalized["audience_action"] = {
            key: value
            for key, value in {
                "action": action,
                "summary": cleaned.pop("action_summary", None),
                "prop": cleaned.pop("prop", None),
                "summoned_actor": cleaned.pop("summoned_actor", None),
            }.items()
            if value is not None
        }

    backend_metadata = {
        key: cleaned.pop(key)
        for key in (
            "backend_name",
            "model_id",
            "load_status",
            "director_mode",
            "tts_backend",
            "voice_mode",
            "voice_name",
        )
        if key in cleaned
    }
    if backend_metadata:
        normalized["backend_metadata"] = backend_metadata

    if "reason_summary" in cleaned:
        normalized["summary"] = cleaned.pop("reason_summary")
    elif "action_summary" in cleaned:
        normalized["summary"] = cleaned.pop("action_summary")
    elif "error_summary" in cleaned:
        normalized["summary"] = cleaned.pop("error_summary")

    for key, value in cleaned.items():
        if key not in normalized:
            normalized[key] = value

    if "summary" not in normalized:
        normalized["summary"] = _event_summary(normalized)
    return {key: normalized[key] for key in _STANDARD_EVENT_KEYS if key in normalized} | {
        key: value for key, value in normalized.items() if key not in _STANDARD_EVENT_KEYS
    }


def _event_summary(event: dict[str, Any]) -> str:
    """
    Generates a short human-readable description of a trace event.

    Used in summaries and logs to quickly understand what happened
    without inspecting full event data.
    """
    event_type = str(event.get("event_type") or "event")
    if event_type == "show_created":
        backend = _backend_label(event)
        return f"Show created with {backend}."
    if event_type == "premise_cast_resolved":
        title = event.get("resolved_show_title") or event.get("show_title") or "?"
        source = event.get("premise_cast_source") or "?"
        fb = event.get("cast_fallback_used")
        n = len(event.get("cast_attempts") or []) if isinstance(event.get("cast_attempts"), list) else 0
        return f"Premise cast: title={title!r} source={source} fallback={fb} attempts={n}."
    if event_type == "actors_created":
        return f"{event.get('actor_count', 'Actors')} actors created."
    if event_type == "director_decision":
        decision = event.get("director_decision") if isinstance(event.get("director_decision"), dict) else {}
        beat_type = decision.get("beat_type") if isinstance(decision, dict) else None
        speaker = event.get("speaker") or "next actor"
        reason = decision.get("reason_summary") if isinstance(decision, dict) else None
        return f"Director chose {speaker} for {beat_type or 'a beat'}" + (f": {reason}" if reason else ".")
    if event_type in {"actor_response", "backend_error", "backend_warmup"}:
        status = event.get("validation_status") or "status unknown"
        fallback = " with fallback" if event.get("fallback_used") else ""
        return f"Backend returned {status}{fallback}."
    if event_type in {"actor_fallback", "director_fallback"}:
        reason = event.get("fallback_reason") or event.get("validation_status") or "fallback used"
        return f"Fallback used: {reason}."
    if event_type == "actor_intent":
        return str(event.get("actor_intent") or "Actor intent recorded.")
    if event_type == "actor_memory_update":
        return str(event.get("actor_memory_update") or "Actor memory updated.")
    if event_type == "tool_requested":
        tool = _tool_label(event.get("tool_request"))
        return f"Tool requested: {tool}."
    if event_type == "tool_executed":
        tool = _tool_label(event.get("tool_request"))
        return f"Tool executed: {tool}."
    if event_type == "tool_result":
        result = event.get("tool_result") if isinstance(event.get("tool_result"), dict) else {}
        return str(result.get("result") or f"Tool result from {result.get('tool_name', 'tool')}.")
    if event_type == "audience_action":
        action = event.get("audience_action") if isinstance(event.get("audience_action"), dict) else {}
        return str(action.get("summary") or action.get("action") or "Audience action recorded.")
    if event_type in {"tts_requested", "tts_generated", "tts_fallback"}:
        metadata = event.get("backend_metadata") if isinstance(event.get("backend_metadata"), dict) else {}
        return f"TTS event via {metadata.get('tts_backend', 'voice backend')}."
    if event_type == "scene_completed":
        return str(event.get("summary") or "Scene completed.")
    return str(event.get("summary") or event_type.replace("_", " ").title())


def sanitize_value(value: Any) -> Any:
    """
    Recursively sanitizes values to remove sensitive or noisy information.

    - Removes secrets, tokens, and credentials
    - Redacts file paths
    - Trims long strings
    - Cleans nested structures (dicts, lists, tuples)
    """
    if isinstance(value, dict):
        return {
            str(key): sanitize_value(item)
            for key, item in value.items()
            if item is not None and not _is_sensitive_key(str(key))
        }
    if isinstance(value, list):
        return [sanitize_value(item) for item in value]
    if isinstance(value, tuple):
        return [sanitize_value(item) for item in value]
    if isinstance(value, str):
        return _sanitize_text(value)
    return value


def _sanitize_text(value: str) -> str:
    """
    Cleans and redacts sensitive information from a string.

    Removes:
    - Extra whitespace
    - Stack traces
    - Environment secrets
    - API keys and credentials
    - Private file system paths

    Also truncates long text for safe logging.
    """
    text = " ".join(value.split())
    if "Traceback (most recent call last)" in text:
        text = "Error details redacted"
    for key, secret in os.environ.items():
        if not _is_sensitive_key(key) or not secret:
            continue
        text = text.replace(secret, "[redacted]")
    text = re.sub(
        r"\b([A-Z][A-Z0-9_]*(?:TOKEN|SECRET|PASSWORD|CREDENTIAL|API_KEY|ACCESS_KEY)[A-Z0-9_]*)\s*=\s*\S+",
        r"\1=[redacted]",
        text,
    )
    for pattern in _PRIVATE_PATH_PATTERNS:
        text = pattern.sub("[redacted-path]", text)
    return text[:_MAX_TEXT_CHARS].rstrip()


def _backend_label(event: dict[str, Any]) -> str:
    """
    Builds a readable label for backend configuration.

    Combines backend name and model ID (if available) to help identify
    which model/engine produced an event.
    """
    metadata = event.get("backend_metadata") if isinstance(event.get("backend_metadata"), dict) else {}
    backend = metadata.get("backend_name") or metadata.get("director_mode") or event.get("backend_name") or "backend"
    model = metadata.get("model_id") or event.get("model_id")
    return f"{backend} ({model})" if model else str(backend)


def _tool_label(tool_request: Any) -> str:
    """
    Extracts a readable tool name from a tool request structure.

    Used in logs and summaries for tool-related events.
    """
    if isinstance(tool_request, dict):
        return str(tool_request.get("tool_name") or "tool")
    return "tool"


def _is_sensitive_key(key: str) -> bool:
    """
    Determines whether a key name contains sensitive information.

    Used to prevent logging of secrets like:
    - tokens
    - passwords
    - API keys
    - credentials
    """
    lowered = key.lower()
    return any(marker in lowered for marker in ("token", "secret", "password", "credential", "api_key", "access_key"))


def _model_id_for_mode(mode: str | None) -> str | None:
    """
    Resolves model identifiers based on backend mode.

    Maps execution modes (e.g., openbmb, hf_api) to their corresponding
    model IDs for trace metadata.
    """
    if mode == "openbmb":
        return os.getenv("OPENBMB_MODEL_ID", "openbmb/MiniCPM5-1B")
    if mode == "hf_api":
        return os.getenv("HF_API_MODEL_ID", "Qwen/Qwen3-4B-Instruct-2507:nscale")
    return None