from __future__ import annotations

import asyncio
import os
from typing import Optional

import openai
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

router = APIRouter(prefix="/dataset")

# ── Config ─────────────────────────────────────────────────────────────────────

MAX_DEPTH       = 3      # max recursion levels
MAX_ROWS        = 10     # cap how many flagged rows we trace (LLM calls are expensive)
SNIPPET_LEN     = 600    # chars sent to LLM per text
PARAPHRASE_THRESHOLD = 0.55
TRANSLATION_THRESHOLD = 0.55


# ── OpenAI client (lazy, picks up key from env) ────────────────────────────────

_client: Optional[openai.AsyncOpenAI] = None

def _get_client() -> openai.AsyncOpenAI:
    global _client
    if _client is None:
        api_key = os.environ.get("OPENAI_API_KEY")
        if not api_key:
            raise RuntimeError("Liability chain tracing is unavailable.")
        _client = openai.AsyncOpenAI(api_key=api_key)
    return _client


# ── LLM helpers ────────────────────────────────────────────────────────────────

_SYS = (
    "You are a forensic text analyst. Answer ONLY with valid JSON matching the "
    "schema in the user message. No markdown, no explanation."
)

async def _llm_json(prompt: str) -> dict:
    client = _get_client()
    resp = await client.chat.completions.create(
        model="gpt-4o-mini",
        temperature=0,
        response_format={"type": "json_object"},
        messages=[
            {"role": "system", "content": _SYS},
            {"role": "user",   "content": prompt},
        ],
    )
    import json
    return json.loads(resp.choices[0].message.content)


async def _detect_paraphrase(text: str) -> dict:
    """
    Returns {is_paraphrase: bool, confidence: float, source_hint: str|null}
    """
    snippet = text[:SNIPPET_LEN]
    prompt = f"""Analyze this text and determine if it is likely a paraphrase of a well-known source
(e.g. Wikipedia, a news article, a book, academic paper, or public document).

Text:
\"\"\"
{snippet}
\"\"\"

Respond with JSON:
{{
  "is_paraphrase": <true|false>,
  "confidence": <0.0-1.0>,
  "source_hint": <brief description of probable source, or null if unknown>
}}"""
    try:
        return await _llm_json(prompt)
    except Exception as e:
        return {"is_paraphrase": False, "confidence": 0.0, "source_hint": None, "error": str(e)}


async def _detect_translation(text: str) -> dict:
    """
    Returns {is_translation: bool, confidence: float, source_lang: str|null, source_hint: str|null}
    """
    snippet = text[:SNIPPET_LEN]
    prompt = f"""Analyze this text and determine if it reads like a machine translation or human translation
from another language. Look for awkward phrasing, calques, or structural patterns typical of translated text.

Text:
\"\"\"
{snippet}
\"\"\"

Respond with JSON:
{{
  "is_translation": <true|false>,
  "confidence": <0.0-1.0>,
  "source_lang": <ISO 639-1 code of the likely source language, or null>,
  "source_hint": <brief note on translation origin, or null>
}}"""
    try:
        return await _llm_json(prompt)
    except Exception as e:
        return {"is_translation": False, "confidence": 0.0, "source_lang": None, "source_hint": None, "error": str(e)}


# ── Recursive chain tracer ─────────────────────────────────────────────────────

async def _trace_chain(
    text: str,
    synthetic_score: float,
    depth: int = 0,
) -> list[dict]:
    """
    Recursively builds a liability chain for one text snippet.
    Returns a list of chain nodes ordered from shallowest to deepest.
    """
    nodes: list[dict] = []

    # Depth 0 is always the synthetic detection (already done upstream)
    if depth == 0:
        nodes.append({
            "depth": 0,
            "type": "synthetic",
            "score": round(synthetic_score, 4),
            "confidence": "high" if synthetic_score >= 0.9 else "medium",
            "source_hint": None,
        })

    if depth >= MAX_DEPTH:
        return nodes

    # Run paraphrase and translation checks in parallel
    para_task = asyncio.create_task(_detect_paraphrase(text))
    trans_task = asyncio.create_task(_detect_translation(text))
    para, trans = await asyncio.gather(para_task, trans_task)

    if para.get("is_paraphrase") and para.get("confidence", 0) >= PARAPHRASE_THRESHOLD:
        nodes.append({
            "depth": depth + 1,
            "type": "paraphrase",
            "score": round(para["confidence"], 4),
            "confidence": "high" if para["confidence"] >= 0.75 else "medium",
            "source_hint": para.get("source_hint"),
        })
        # Recurse: a paraphrase can itself be a translation
        deeper = await _trace_chain(text, synthetic_score, depth=depth + 1)
        # Only add translation nodes from recursion (skip duplicated synthetic root)
        nodes.extend([n for n in deeper if n["depth"] > depth + 1])

    if trans.get("is_translation") and trans.get("confidence", 0) >= TRANSLATION_THRESHOLD:
        # Only add translation node if not already at this depth from paraphrase branch
        already_at_depth = any(n["depth"] == depth + 1 for n in nodes)
        new_depth = (depth + 2) if already_at_depth else (depth + 1)
        nodes.append({
            "depth": new_depth,
            "type": "translation",
            "score": round(trans["confidence"], 4),
            "confidence": "high" if trans["confidence"] >= 0.75 else "medium",
            "source_lang": trans.get("source_lang"),
            "source_hint": trans.get("source_hint"),
        })

    return nodes


def _liability_score(chain: list[dict]) -> float:
    """
    Composite liability score (0–1) based on chain depth and node confidence scores.
    Each deeper node multiplies liability — synthetic alone = its score,
    each additional layer adds compounding risk.
    """
    if not chain:
        return 0.0
    base = chain[0]["score"]
    multiplier = 1.0
    for node in chain[1:]:
        multiplier *= (1 + 0.3 * node["score"])
    raw = base * multiplier
    return round(min(raw, 1.0), 4)


# ── Core analysis ──────────────────────────────────────────────────────────────

async def run_liability_chain_analysis(
    flagged_rows: list[dict],   # from synthetic result: [{row_idx, column, synthetic_score, confidence}, ...]
    all_rows: list[dict],       # full dataset rows for text lookup
    max_rows: int = MAX_ROWS,
) -> dict:
    # Take the top-scoring flagged rows
    candidates = sorted(flagged_rows, key=lambda r: r["synthetic_score"], reverse=True)[:max_rows]

    chains: list[dict] = []

    async def process_one(flagged: dict) -> dict:
        row_idx = flagged["row_idx"]
        column  = flagged["column"]
        score   = flagged["synthetic_score"]

        # Look up the actual text
        text = ""
        if row_idx < len(all_rows):
            val = all_rows[row_idx].get(column, "")
            if isinstance(val, str):
                text = val.strip()

        if not text or len(text) < 20:
            return {
                "row_idx":       row_idx,
                "column":        column,
                "synthetic_score": score,
                "chain":         [{"depth":0,"type":"synthetic","score":score,"confidence":flagged["confidence"],"source_hint":None}],
                "max_depth":     0,
                "liability_score": round(score, 4),
                "skipped":       True,
            }

        chain = await _trace_chain(text, score, depth=0)
        return {
            "row_idx":         row_idx,
            "column":          column,
            "synthetic_score": score,
            "chain":           chain,
            "max_depth":       max(n["depth"] for n in chain) if chain else 0,
            "liability_score": _liability_score(chain),
            "text_snippet":    text[:200],
        }

    tasks = [asyncio.create_task(process_one(f)) for f in candidates]
    results = await asyncio.gather(*tasks)
    chains = list(results)

    # Sort by liability score desc
    chains.sort(key=lambda c: c["liability_score"], reverse=True)

    avg_liability = round(sum(c["liability_score"] for c in chains) / len(chains), 4) if chains else 0.0
    max_depth     = max((c["max_depth"] for c in chains), default=0)
    deep_chains   = sum(1 for c in chains if c["max_depth"] >= 2)

    return {
        "rows_traced":    len(chains),
        "avg_liability":  avg_liability,
        "max_depth_seen": max_depth,
        "deep_chains":    deep_chains,
        "chains":         chains,
    }


# ── Request model ──────────────────────────────────────────────────────────────

class LiabilityChainRequest(BaseModel):
    flagged_rows: list[dict]   # [{row_idx, column, synthetic_score, confidence}]
    rows:         list[dict]   # full dataset rows
    max_rows:     int = MAX_ROWS


# ── Route ──────────────────────────────────────────────────────────────────────

@router.post("/liability-chain")
async def liability_chain(req: LiabilityChainRequest):
    if not req.flagged_rows:
        raise HTTPException(status_code=400, detail="No flagged rows provided")
    try:
        return await run_liability_chain_analysis(
            req.flagged_rows,
            req.rows,
            max_rows=req.max_rows,
        )
    except RuntimeError as e:
        raise HTTPException(status_code=503, detail=str(e))