from __future__ import annotations

from hashlib import sha256
import json
from pathlib import Path
from typing import Any, Dict, List
import time

from armorer.security.policy import PolicyGuard


_SECRET_KEY_TOKENS = ("key", "secret", "token", "password", "pass", "auth")


def _base_dir() -> Path:
    root = Path.home() / ".armorer" / "training_data"
    (root / "raw").mkdir(parents=True, exist_ok=True)
    (root / "exports").mkdir(parents=True, exist_ok=True)
    return root


def _is_secret_key(key: str) -> bool:
    lower = str(key or "").lower()
    return any(token in lower for token in _SECRET_KEY_TOKENS)


def _sanitize_value(value: Any, guard: PolicyGuard, key_hint: str = "") -> Any:
    if isinstance(value, dict):
        out: Dict[str, Any] = {}
        for k, v in value.items():
            if _is_secret_key(str(k)):
                out[str(k)] = "[REDACTED]"
            else:
                out[str(k)] = _sanitize_value(v, guard, key_hint=str(k))
        return out
    if isinstance(value, list):
        return [_sanitize_value(item, guard, key_hint=key_hint) for item in value]
    if isinstance(value, str):
        if _is_secret_key(key_hint):
            return "[REDACTED]"
        return guard.sanitize_text(value)
    if value is None or isinstance(value, (int, float, bool)):
        return value
    if isinstance(value, Path):
        return guard.sanitize_text(str(value))
    if hasattr(value, "model_dump"):
        try:
            return _sanitize_value(value.model_dump(), guard, key_hint=key_hint)
        except Exception:
            return guard.sanitize_text(str(value))
    if hasattr(value, "dict") and callable(getattr(value, "dict")):
        try:
            return _sanitize_value(value.dict(), guard, key_hint=key_hint)
        except Exception:
            return guard.sanitize_text(str(value))
    if hasattr(value, "__dict__"):
        try:
            return _sanitize_value(vars(value), guard, key_hint=key_hint)
        except Exception:
            return guard.sanitize_text(str(value))
    return guard.sanitize_text(str(value))


def _append_jsonl(path: Path, record: Dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as fh:
        fh.write(json.dumps(record, ensure_ascii=True) + "\n")


def _system_prompt_digest(text: str) -> str:
    return sha256((text or "").encode("utf-8")).hexdigest()[:16]


def _trace_to_sft(trace: Dict[str, Any]) -> Dict[str, Any] | None:
    final = str(trace.get("final_response") or "").strip()
    messages = trace.get("messages") or []
    if not final or not messages:
        return None
    if trace.get("outcome") != "success":
        return None
        
    # The messages array now contains the full history including tool calls
    sft_messages = []
    for msg in messages:
        # Copy exactly so tool_calls, tool_call_id, etc. are preserved
        sft_msg = dict(msg)
        sft_messages.append(sft_msg)
        
    return {
        "trace_id": trace.get("trace_id"),
        "messages": sft_messages,
        "metadata": {
            "provider": trace.get("provider"),
            "model": trace.get("model"),
            "task_type": trace.get("task_type"),
            "attempt_count": len(trace.get("attempts") or []),
            "system_prompt_digest": trace.get("system_prompt_digest"),
        },
    }


def _trace_to_trajectory(trace: Dict[str, Any]) -> Dict[str, Any]:
    return {
        "trace_id": trace.get("trace_id"),
        "timestamp": trace.get("timestamp"),
        "prompt": trace.get("messages"),
        "final_response": trace.get("final_response"),
        "provider": trace.get("provider"),
        "model": trace.get("model"),
        "task_type": trace.get("task_type"),
        "outcome": trace.get("outcome"),
        "failure_labels": trace.get("failure_labels") or [],
        "attempts": trace.get("attempts") or [],
        "tool_names_exposed": trace.get("tool_names_exposed") or [],
        "system_prompt_digest": trace.get("system_prompt_digest"),
    }


def record_training_trace(
    *,
    trace_id: str,
    session_id: str,
    messages: List[Dict[str, Any]],
    provider: str,
    model: str,
    system_prompt: str,
    task_type: str,
    tool_names_exposed: List[str],
    attempts: List[Dict[str, Any]],
    final_response: str,
    outcome: str,
    failure_labels: List[str],
    guard: PolicyGuard,
) -> Dict[str, Any]:
    root = _base_dir()
    timestamp = time.time()
    trace = {
        "trace_id": trace_id,
        "session_id": session_id,
        "timestamp": timestamp,
        "messages": _sanitize_value(messages, guard),
        "provider": provider,
        "model": model,
        "task_type": task_type,
        "tool_names_exposed": _sanitize_value(tool_names_exposed, guard),
        "system_prompt_digest": _system_prompt_digest(system_prompt),
        "attempts": _sanitize_value(attempts, guard),
        "final_response": _sanitize_value(final_response, guard),
        "outcome": outcome,
        "failure_labels": _sanitize_value(list(dict.fromkeys(failure_labels)), guard),
    }

    date_dir = time.strftime("%Y-%m-%d", time.localtime(timestamp))
    raw_path = root / "raw" / date_dir / f"{trace_id}.json"
    raw_path.parent.mkdir(parents=True, exist_ok=True)
    raw_path.write_text(json.dumps(trace, indent=2, ensure_ascii=True), encoding="utf-8")

    trajectory = _trace_to_trajectory(trace)
    _append_jsonl(root / "exports" / "trajectories.jsonl", trajectory)

    sft = _trace_to_sft(trace)
    if sft:
        _append_jsonl(root / "exports" / "sft.jsonl", sft)

    return {
        "raw_path": str(raw_path),
        "trajectory_export": str(root / "exports" / "trajectories.jsonl"),
        "sft_export": str(root / "exports" / "sft.jsonl"),
    }
