from __future__ import annotations

import json
import os
import sys
import time
from typing import Any, Dict, List

import requests
from openai import OpenAI
from pydantic import ValidationError

from models import EraseOpsAction

API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
HF_TOKEN = os.getenv("HF_TOKEN")
TASK_NAME = os.getenv("TASK_NAME", "clean_sweep")
ENV_URL = os.getenv("ENV_URL", "http://localhost:7860")
MAX_TOTAL_RUNTIME_SECONDS = int(os.getenv("MAX_TOTAL_RUNTIME_SECONDS", "1100"))
MAX_TASK_RUNTIME_SECONDS = int(os.getenv("MAX_TASK_RUNTIME_SECONDS", "260"))

SYSTEM_PROMPT = (
    "You are a Data Privacy Officer agent managing GDPR Right-to-Erasure requests.\n"
    "At each step you receive a JSON observation showing all data records across enterprise systems.\n"
    "Your goal: delete all PII records belonging to the target user_id in the request.\n\n"
    "Rules you must follow:\n"
    "1. NEVER delete or anonymize records where legal_hold=True - escalate these instead\n"
    "2. NEVER delete records belonging to a different user_id than in the request\n"
    "3. Always delete dependency records BEFORE the records that depend on them\n"
    "4. Use 'inspect' action first if you are unsure about a record's dependencies\n"
    "5. Use 'anonymize' only for derived or anonymized type records, never for PII\n\n"
    "Respond with ONLY a valid JSON object matching this schema:\n"
    "{\n"
    "  'action_type': 'inspect' | 'delete' | 'anonymize' | 'escalate' | 'noop',\n"
    "  'record_id': 'rec_001',\n"
    "  'justification': 'brief reason'\n"
    "}"
)

client = OpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL)


def _log_start(task: str, model: str) -> None:
    print(f"[START] task={task} env=gdpr-eraseops model={model}", flush=True)


def _log_step(step: int, action_json: str, reward: float, done: bool, error: str | None) -> None:
    error_text = "null" if error in (None, "", "None") else error
    print(
        f"[STEP] step={step} action={action_json} reward={reward:.2f} done={str(done).lower()} error={error_text}",
        flush=True,
    )


def _log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    rewards_str = ",".join(f"{value:.2f}" for value in rewards)
    print(f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_str}", flush=True)


def _safe_action() -> Dict[str, Any]:
    return {"action_type": "noop", "record_id": None, "justification": "safe fallback"}


def _candidate_record(observation: Dict[str, Any]) -> str | None:
    request_user = observation.get("request", {}).get("user_id")
    records = observation.get("records", [])

    for record in records:
        if record.get("status") != "active":
            continue
        if record.get("legal_hold") and record.get("user_id") == request_user:
            return str(record.get("record_id"))

    for record in records:
        if record.get("status") != "active":
            continue
        if record.get("user_id") != request_user:
            continue
        if record.get("data_type") != "PII":
            continue
        dependencies = record.get("dependencies", [])
        deps_deleted = True
        for dep_id in dependencies:
            dep = next((r for r in records if r.get("record_id") == dep_id), None)
            if dep is None or dep.get("status") != "deleted":
                deps_deleted = False
                break
        if deps_deleted:
            return str(record.get("record_id"))

    return None


def _heuristic_action(observation: Dict[str, Any]) -> Dict[str, Any]:
    request_user = observation.get("request", {}).get("user_id")
    records = observation.get("records", [])
    step = int(observation.get("step", 0) or 0)

    if step < min(3, len(records)):
        rec = records[step]
        return {
            "action_type": "inspect",
            "record_id": rec.get("record_id"),
            "justification": "progressive reconnaissance",
        }

    for record in records:
        if (
            record.get("status") == "active"
            and record.get("user_id") == request_user
            and record.get("legal_hold") is True
        ):
            return {
                "action_type": "escalate",
                "record_id": record.get("record_id"),
                "justification": "legal hold discovered",
            }

    candidate = _candidate_record(observation)
    if candidate:
        return {
            "action_type": "delete",
            "record_id": candidate,
            "justification": "delete active in-scope PII with dependencies satisfied",
        }

    return {"action_type": "noop", "record_id": None, "justification": "no safe move"}


def _llm_action(observation: Dict[str, Any]) -> Dict[str, Any]:
    if not HF_TOKEN:
        return _heuristic_action(observation)

    user_prompt = (
        f"{json.dumps(observation, separators=(',', ':'), ensure_ascii=False)}\n"
        "Choose your next action. Respond with JSON only."
    )

    try:
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.0,
            max_tokens=220,
        )
        raw = (response.choices[0].message.content or "").strip()
        if raw.startswith("```"):
            raw = raw.replace("```json", "").replace("```", "").strip()
        parsed = json.loads(raw)
        return {
            "action_type": parsed.get("action_type", "noop"),
            "record_id": parsed.get("record_id"),
            "justification": parsed.get("justification"),
        }
    except Exception:
        return _heuristic_action(observation)


def _validated_action(raw_action: Dict[str, Any]) -> Dict[str, Any]:
    try:
        return EraseOpsAction(**raw_action).model_dump()
    except ValidationError:
        return _safe_action()


def _run_task(task: str) -> float:
    _log_start(task, MODEL_NAME)
    task_started_at = time.monotonic()

    reset_resp = requests.post(
        f"{ENV_URL}/reset",
        params={"task_name": task, "seed": 42},
        timeout=20,
    )
    reset_resp.raise_for_status()
    observation = reset_resp.json()

    rewards: List[float] = []
    done = False
    step = 0

    max_steps = int(observation.get("step_budget", 30) or 30)

    while not done and step < max_steps:
        if time.monotonic() - task_started_at > MAX_TASK_RUNTIME_SECONDS:
            print(f"Runtime guard: task '{task}' timed out", file=sys.stderr, flush=True)
            done = True
            break

        step += 1
        raw_action = _llm_action(observation)
        action = _validated_action(raw_action)
        action_json = json.dumps(action, separators=(",", ":"), ensure_ascii=False)

        error_text: str | None = None
        reward_total = 0.0

        try:
            step_resp = requests.post(f"{ENV_URL}/step", json=action, timeout=20)
            step_resp.raise_for_status()
            payload = step_resp.json()
            observation = payload.get("observation", {})
            reward_total = float(payload.get("reward", {}).get("total", 0.0) or 0.0)
            done = bool(payload.get("done", False))
            info = payload.get("info", {})
            error_text = info.get("error") if isinstance(info, dict) else None
        except Exception as exc:
            done = True
            reward_total = -1.0
            error_text = str(exc)

        rewards.append(reward_total)
        _log_step(step=step, action_json=action_json, reward=reward_total, done=done, error=error_text)

    grade_resp = requests.post(f"{ENV_URL}/grade", params={"task_name": task}, timeout=20)
    grade_resp.raise_for_status()
    score = float(grade_resp.json().get("score", 0.0) or 0.0)

    success = score >= 0.5
    _log_end(success=success, steps=step, score=score, rewards=rewards)
    return score


def main() -> None:
    """Run the full baseline across all tasks."""
    tasks = ["clean_sweep", "tangled_lake", "legal_minefield", "cascade_collapse"]
    run_started_at = time.monotonic()

    for task in tasks:
        if time.monotonic() - run_started_at > MAX_TOTAL_RUNTIME_SECONDS:
            print("Runtime guard: stopping before timeout", file=sys.stderr, flush=True)
            break
        _run_task(task)


if __name__ == "__main__":
    main()