"""
Stage 3 - Constrained Forensic Narrative Generation.
Uses OpenRouter (Gemma 4 Vision free tier) with strict grounding to structured findings.
Includes post-generation validation to prevent hallucination.
"""

import os
import base64
import json
import httpx
from typing import Optional, Dict, Any, List
from dataclasses import dataclass

# ── Config ──────────────────────────────────────────────
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1/chat/completions"

# Primary: Known working free vision model
PRIMARY_MODEL = "inclusionai/ling-3.0-flash-vl:free"
# Fallback 1: High-throughput free reasoning model
FALLBACK_MODEL_1 = "nvidia/nemotron-3.5-lightning:free"
# Fallback 2: Gemma 4 free instruction model
FALLBACK_MODEL_2 = "google/gemma-4-31b-it:free"

OPENROUTER_TIMEOUT = 60.0  # seconds


# ── Constrained Prompt Builder ──────────────────────────

def build_constrained_prompt(findings: Dict[str, Any]) -> str:
    """
    Build a prompt that constrains the VLM to ONLY use the structured findings.
    No free-form image analysis allowed.
    """
    ensemble = findings.get("ensemble", {})
    artifacts = findings.get("artifacts", [])
    metadata = findings.get("metadata", {})
    uncertainty = findings.get("uncertainty", {})

    ensemble_label = ensemble.get("ensemble_label", "UNCERTAIN")
    ensemble_conf = ensemble.get("ensemble_confidence", 0.0)
    disagreement = ensemble.get("ensemble_disagreement", 0.0)
    votes = ensemble.get("votes", [])
    primary_model = ensemble.get("primary_model", "unknown")

    # Format votes for prompt
    vote_lines = []
    for v in votes:
        if v.get("weight", 0) > 0:
            vote_lines.append(
                f"- {v['model_name']}: {v['label']} (P(FAKE)={v['raw_score']:.3f}, weight={v['weight']:.2f})"
            )
    votes_text = "\n".join(vote_lines) if vote_lines else "No valid detector votes."

    # Format artifacts
    artifact_lines = []
    for a in artifacts:
        if isinstance(a, dict):
            art_type = a.get("type", "unknown")
            desc = a.get("description", "")
            sev = a.get("severity", 0)
            artifact_lines.append(f"- {art_type}: {desc} (severity={sev:.2f})")
        else:
            artifact_lines.append(f"- {a}")
    artifacts_text = "\n".join(artifact_lines) if artifact_lines else "No specific artifacts detected."

    # Format metadata
    meta_summary = metadata.get("summary", "No metadata analysis.")
    meta_weight = metadata.get("weight", "minor")
    meta_absent = metadata.get("metadata_absent", False)

    # Uncertainty info
    uncertain = uncertainty.get("ensemble_disagreement", 0) > 0.25 or ensemble_label == "UNCERTAIN"
    oos = uncertainty.get("oos_flag", False)

    prompt = f"""You are a forensic image analyst. Write a concise forensic report using ONLY the structured findings below.

STRUCTURED FINDINGS (DO NOT DEVIATE FROM THESE):

ENSEMBLE CLASSIFICATION:
- Verdict: {ensemble_label}
- Confidence: {ensemble_conf:.1%}
- Primary Detector: {primary_model}
- Detector Votes:
{votes_text}
- Ensemble Disagreement (std of P(FAKE)): {disagreement:.3f}
- Uncertainty Flag: {"YES - detectors disagreed significantly" if uncertain else "NO"}
- Out-of-Distribution Flag: {"YES" if oos else "NO"}

NAMED ARTIFACTS DETECTED:
{artifacts_text}

METADATA ANALYSIS (Weight: {meta_weight.upper()} - Corroborating Only):
- Summary: {meta_summary}
- Metadata Absent: {"Yes" if meta_absent else "No"}

RULES - YOU MUST FOLLOW THESE:
1. ONLY reference findings explicitly listed above. Do NOT infer visual details not in the findings.
2. If Uncertainty Flag is YES, you MUST state: "Low confidence — detectors disagreed significantly."
3. If Out-of-Distribution Flag is YES, you MUST state: "Input may be outside detector training distribution."
4. Do NOT claim lighting inconsistencies, shadow analysis, facial asymmetry, or skin texture analysis unless explicitly in the artifacts list.
5. Do NOT make up specific locations (e.g., "left eye", "jawline") unless in artifacts.
6. Metadata weight is MINOR - never let metadata alone drive conclusions.
7. Structure your report as:
   A. VERDICT & CONFIDENCE
   B. KEY EVIDENCE (cite specific detector votes and artifacts)
   C. UNCERTAINTY & LIMITATIONS
   D. RECOMMENDATION (one of: FLAG FOR REVIEW / MONITOR / LIKELY AUTHENTIC)
8. Keep under 250 words. Be direct and forensically precise.

REPORT:"""

    return prompt


# ── OpenRouter Client ───────────────────────────────────

async def call_openrouter(
    original_b64: str,
    heatmap_b64: str,
    prompt: str,
    model: str,
) -> Optional[str]:
    """Call OpenRouter API with multimodal input."""
    if not OPENROUTER_API_KEY:
        print("[LLM] No OPENROUTER_API_KEY set")
        return None

    payload = {
        "model": model,
        "messages": [{
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{original_b64}"}
                },
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{heatmap_b64}"}
                }
            ]
        }],
        "max_tokens": 500,
        "temperature": 0.3,
        "top_p": 0.9,
    }

    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://deepguard.local",
        "X-Title": "DeepGuard Forensic Analysis",
    }

    try:
        async with httpx.AsyncClient(timeout=OPENROUTER_TIMEOUT) as client:
            resp = await client.post(OPENROUTER_BASE_URL, json=payload, headers=headers)
            
            if resp.status_code == 429:
                print("[LLM] Rate limited (429)")
                return None
            
            if resp.status_code != 200:
                print(f"[LLM] OpenRouter error {resp.status_code}: {resp.text}")
                return None
            
            data = resp.json()
            content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
            return content.strip() if content else None
            
    except httpx.TimeoutException:
        print("[LLM] OpenRouter timeout")
        return None
    except Exception as e:
        print(f"[LLM] OpenRouter call failed: {type(e).__name__}: {e}")
        return None


# ── Response Validation ──────────────────────────────────

def validate_narrative(narrative: str, findings: Dict[str, Any]) -> str:
    """
    Post-process validation: check narrative doesn't hallucinate claims.
    Returns cleaned narrative with hallucinated claims removed/flagged.
    """
    if not narrative:
        return ""
    
    ensemble = findings.get("ensemble", {})
    artifacts = findings.get("artifacts", [])
    metadata = findings.get("metadata", {})
    
    # Extract allowed terms from findings
    allowed_terms = set()
    
    # From ensemble
    for vote in ensemble.get("votes", []):
        allowed_terms.add(vote.get("model_name", "").lower())
        allowed_terms.add(vote.get("label", "").lower())
    
    # From artifacts
    for a in artifacts:
        if isinstance(a, dict):
            allowed_terms.add(a.get("type", "").lower())
            desc = a.get("description", "").lower()
            allowed_terms.update(desc.split())
    
    # From metadata
    if metadata.get("software"):
        allowed_terms.add(metadata["software"].lower())
    
    # Common forensic terms that are OK
    allowed_terms.update({
        "fake", "real", "authentic", "manipulated", "generated", "synthetic",
        "detector", "ensemble", "confidence", "disagreement", "uncertainty",
        "artifact", "evidence", "analysis", "forensic", "verdict", "recommendation",
        "review", "monitor", "likely", "probability", "model", "vote", "weight",
        "metadata", "exif", "software", "camera", "dimensions", "square",
        "upsampling", "checkerboard", "diffusion", "gan", "blending", "boundary",
        "ela", "compression", "inconsistency", "spectral", "frequency", "heatmap",
        "grad-cam", "activation", "attention", "region", "severity", "flag",
        "low", "high", "moderate", "significant", "limited", "corroborating",
    })
    
    # Check for hallucinated specific claims
    narrative_lower = narrative.lower()
    
    # Patterns that indicate hallucination (specific visual claims not in findings)
    hallucination_patterns = [
        "lighting inconsistency",
        "shadow inconsistency",
        "skin texture",
        "facial asymmetry",
        "eye reflection",
        "jawline",
        "hairline",
        "ear",
        "neck boundary",
        "blending edge",
        "color mismatch",
        "unnatural blink",
        "temporal inconsistency",
        "lip sync",
    ]
    
    flagged_claims = []
    for pattern in hallucination_patterns:
        if pattern in narrative_lower:
            # Check if this pattern is actually in artifacts
            in_artifacts = any(pattern in str(a).lower() for a in artifacts)
            if not in_artifacts:
                flagged_claims.append(pattern)
    
    if flagged_claims:
        print(f"[LLM Validation] Flagged potential hallucinations: {flagged_claims}")
        # Append disclaimer
        narrative += f"\n\n[VALIDATION NOTE: The following claims were not found in structured findings and may be hallucinated: {', '.join(flagged_claims)}]"
    
    return narrative


# ── Template Fallback ───────────────────────────────────

def generate_template_narrative(findings: Dict[str, Any]) -> Dict[str, Any]:
    """Rule-based narrative when LLM unavailable."""
    ensemble = findings.get("ensemble", {})
    label = ensemble.get("ensemble_label", "UNCERTAIN")
    confidence = ensemble.get("ensemble_confidence", 0.0)
    votes = ensemble.get("votes", [])
    artifacts = findings.get("artifacts", [])
    metadata = findings.get("metadata", {})
    uncertainty = findings.get("uncertainty", {})
    
    pct = confidence * 100
    is_fake = label == "FAKE"
    uncertain = uncertainty.get("ensemble_disagreement", 0) > 0.25 or label == "UNCERTAIN"
    
    # Build evidence summary
    vote_summary = "; ".join([
        f"{v['model_name']}: {v['label']} ({v['raw_score']:.0%})" 
        for v in votes if v.get("weight", 0) > 0
    ])
    
    artifact_summary = ", ".join([
        a.get("type", str(a)) if isinstance(a, dict) else str(a)
        for a in artifacts[:3]
    ]) if artifacts else "none"
    
    if label == "UNCERTAIN" or uncertain:
        narrative = (
            f"The ensemble classifier observed conflicting signals across detectors ({pct:.1f}% estimated manipulation score). "
            f"Detector votes: {vote_summary}. "
            f"Artifacts detected: {artifact_summary}. "
            f"Metadata: {metadata.get('summary', 'N/A')}. "
            f"LOW CONFIDENCE — detectors disagreed significantly. Continuous monitoring and manual corroboration advised."
        )
        recommendation = "MONITOR"
    elif is_fake:
        if confidence > 0.90:
            narrative = (
                f"The ensemble classifier assigned a very high probability ({pct:.1f}%) that this image is AI-generated/manipulated. "
                f"Detector votes: {vote_summary}. "
                f"Artifacts detected: {artifact_summary}. "
                f"Metadata: {metadata.get('summary', 'N/A')}. "
                f"Recommend flagging this content for immediate human expert review."
            )
            recommendation = "FLAG FOR REVIEW"
        elif confidence > 0.70:
            narrative = (
                f"The ensemble classifier reports a {pct:.1f}% probability of manipulation. "
                f"Detector votes: {vote_summary}. "
                f"Artifacts detected: {artifact_summary}. "
                f"Evidence is significant but not conclusive. Additional analysis recommended."
            )
            recommendation = "FLAG FOR REVIEW"
        else:
            narrative = (
                f"The ensemble classifier reports a moderate probability ({pct:.1f}%) of manipulation. "
                f"Detector votes: {vote_summary}. "
                f"Artifacts detected: {artifact_summary}. "
                f"Exercise caution — this result warrants monitoring."
            )
            recommendation = "MONITOR"
    else:
        narrative = (
            f"The ensemble classifier assigned a low probability ({pct:.1f}%) of manipulation. "
            f"Detector votes: {vote_summary}. "
            f"Artifacts detected: {artifact_summary}. "
            f"Metadata: {metadata.get('summary', 'N/A')}. "
            f"The image appears consistent with authentic capture."
        )
        recommendation = "LIKELY AUTHENTIC"
    
    return {
        "prediction": label,
        "confidence": round(confidence, 4),
        "key_artifacts": [a.get("type", str(a)) if isinstance(a, dict) else str(a) for a in artifacts[:8]],
        "forensic_narrative": narrative,
        "recommendation": recommendation,
        "llm_used": False,
    }


# ── Main Entry Point ─────────────────────────────────────

async def generate_forensic_report(
    original_bytes: bytes,
    heatmap_bytes: bytes,
    ensemble_result: Any,  # EnsembleResult from stage1_ensemble
    artifacts: List[Dict],
    metadata: Dict,
    mode: str = "deepfake",
    use_llm: bool = True,
) -> Dict[str, Any]:
    """
    Orchestrate Stage 3:
      1. Try OpenRouter (Llama 3.2 Vision) — unless use_llm is False
      2. Fallback 1: Phi 3.5 Vision
      3. Fallback 2: OpenRouter auto-router
      4. Fallback to template narrative
    
    All inputs are structured findings - no raw image analysis by LLM.
    """
    # Fast mode - skip LLM entirely
    if not use_llm:
        print("[LLM] Skipped by user — using template narrative (fast mode)")
        findings = {
            "ensemble": ensemble_result.__dict__ if hasattr(ensemble_result, '__dict__') else ensemble_result,
            "artifacts": artifacts,
            "metadata": metadata,
            "uncertainty": {
                "ensemble_disagreement": getattr(ensemble_result, 'ensemble_disagreement', 0),
                "oos_flag": getattr(ensemble_result, 'oos_flag', False),
            }
        }
        return generate_template_narrative(findings)

    # Prepare structured findings for prompt
    findings = {
        "ensemble": ensemble_result.__dict__ if hasattr(ensemble_result, '__dict__') else ensemble_result,
        "artifacts": artifacts,
        "metadata": metadata,
        "uncertainty": {
            "ensemble_disagreement": getattr(ensemble_result, 'ensemble_disagreement', 0),
            "oos_flag": getattr(ensemble_result, 'oos_flag', False),
        }
    }

    # Build constrained prompt
    prompt = build_constrained_prompt(findings)

    # Encode images
    original_b64 = base64.b64encode(original_bytes).decode("utf-8")
    heatmap_b64 = base64.b64encode(heatmap_bytes).decode("utf-8")

    # Try primary model (Llama 3.2 Vision)
    print(f"[LLM] Calling OpenRouter with {PRIMARY_MODEL}...")
    llm_text = await call_openrouter(original_b64, heatmap_b64, prompt, PRIMARY_MODEL)

    # Fallback 1: Phi 3.5 Vision
    if not llm_text or len(llm_text) < 50:
        print(f"[LLM] Primary failed, trying {FALLBACK_MODEL_1}...")
        llm_text = await call_openrouter(original_b64, heatmap_b64, prompt, FALLBACK_MODEL_1)

    # Fallback 2: OpenRouter auto-router
    if not llm_text or len(llm_text) < 50:
        print(f"[LLM] Fallback 1 failed, trying {FALLBACK_MODEL_2}...")
        llm_text = await call_openrouter(original_b64, heatmap_b64, prompt, FALLBACK_MODEL_2)

    # Final fallback to template
    if not llm_text or len(llm_text) < 50:
        print("[LLM] All OpenRouter models failed — using template narrative")
        return generate_template_narrative(findings)

    # Validate narrative against findings
    print(f"[LLM] Response received ({len(llm_text)} chars)")
    validated_text = validate_narrative(llm_text, findings)

    # Parse response for structured output
    return parse_response(validated_text, findings)


def parse_response(text: str, findings: Dict[str, Any]) -> Dict[str, Any]:
    """Extract structured fields from validated LLM response."""
    text_lower = text.lower()
    ensemble = findings.get("ensemble", {})
    
    # Classification from ensemble (not LLM)
    label = ensemble.get("ensemble_label", "UNCERTAIN")
    confidence = ensemble.get("ensemble_confidence", 0.0)
    
    # Confidence modifiers from text
    conf_map = {"very high": 0.95, "high": 0.85, "medium": 0.65, "moderate": 0.65, "low": 0.40}
    reported_conf = confidence
    for kw, cv in conf_map.items():
        if kw in text_lower:
            reported_conf = cv
            break
    
    # Artifacts from text (but limited to those in findings)
    artifacts = findings.get("artifacts", [])
    detected_artifacts = []
    for a in artifacts:
        if isinstance(a, dict):
            art_type = a.get("type", "")
            if art_type.lower() in text_lower:
                detected_artifacts.append(art_type)
        elif str(a).lower() in text_lower:
            detected_artifacts.append(str(a))
    
    # Recommendation - strictly grounded in ensemble prediction label
    is_fake = "fake" in label.lower() or "ai" in label.lower()
    if is_fake:
        if "monitor" in text_lower or reported_conf < 0.70:
            recommendation = "MONITOR"
        else:
            recommendation = "FLAG FOR REVIEW"
    else:
        if "flag for review" in text_lower:
            recommendation = "MONITOR"  # Cautious review for real images
        elif "monitor" in text_lower:
            recommendation = "MONITOR"
        else:
            recommendation = "LIKELY AUTHENTIC"
    
    return {
        "prediction": label,
        "confidence": round(reported_conf, 4),
        "key_artifacts": detected_artifacts[:8],
        "forensic_narrative": text,
        "recommendation": recommendation,
        "llm_used": True,
    }


# ── Interactive Forensic Chat ───────────────────────────

async def chat_with_forensic_context(
    messages: List[Dict[str, str]],
    forensic_context: Dict[str, Any],
) -> Dict[str, Any]:
    """
    Multi-turn conversation between human investigator and forensic assistant.
    Grounds the discussion in the actual findings while synthesizing investigator leads.
    """
    label = forensic_context.get("label", "UNKNOWN")
    confidence = forensic_context.get("confidence", 0.0)
    mode = forensic_context.get("mode", "ai_image")
    narrative = forensic_context.get("narrative", "")
    artifacts = forensic_context.get("key_artifacts", [])
    votes = forensic_context.get("ensemble_votes", [])
    metadata = forensic_context.get("metadata", {})

    # Format votes
    votes_summary = ", ".join([
        f"{v.get('model_name', 'Detector')}: {v.get('label', '')} ({round(v.get('confidence', 0) * 100)}%)"
        for v in votes if isinstance(v, dict)
    ]) or "Standard ensemble"

    artifacts_summary = ", ".join([str(a) for a in artifacts]) or "None specifically flagged"
    meta_summary = metadata.get("summary", "No EXIF available") if isinstance(metadata, dict) else "N/A"

    system_prompt = f"""You are DeepGuard's Forensic AI Investigator Assistant.
You have just analyzed an image with the following technical findings:
- Mode: {mode.upper()}
- Ensemble Verdict: {label} ({round(confidence * 100)}% confidence)
- Detector Breakdown: {votes_summary}
- Detected Artifacts: {artifacts_summary}
- Metadata Context: {meta_summary}
- Initial Forensic Report: {narrative}

The human investigator is now discussing the case with you. They may provide external context (e.g. source platform, photographer claim, upload compression, generator clues) or ask clarifying questions.
GUIDELINES:
1. Be professional, objective, and scientifically grounded in image forensics.
2. Incorporate the human's newly supplied context (e.g., if they mention Reddit, Midjourney, or camera specs, explain how that interacts with the detector signals).
3. If they ask how to interpret the Grad-CAM++ or frequency signals, give clear, actionable explanations.
4. Do not state assumptions as proven facts — distinguish between hard technical signals and circumstantial human context.
5. Keep responses concise (2-4 clear paragraphs or bullet points)."""

    formatted_messages = [{"role": "system", "content": system_prompt}]
    for m in messages[-10:]:  # Keep last 10 messages for context
        role = m.get("role", "user")
        content = m.get("content", "")
        if role in ["user", "assistant"]:
            formatted_messages.append({"role": role, "content": content})

    payload = {
        "model": PRIMARY_MODEL,
        "messages": formatted_messages,
        "max_tokens": 600,
        "temperature": 0.4,
    }

    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://deepguard.local",
        "X-Title": "DeepGuard Forensic Chat",
    }

    models_to_try = [
        PRIMARY_MODEL,
        "meta-llama/llama-3.3-70b-instruct:free",
        FALLBACK_MODEL_1,
        FALLBACK_MODEL_2,
    ]

    reply_text = None
    used_model = None

    async with httpx.AsyncClient(timeout=45.0) as client:
        for model_name in models_to_try:
            payload["model"] = model_name
            try:
                resp = await client.post(OPENROUTER_BASE_URL, json=payload, headers=headers)
                if resp.status_code == 200:
                    data = resp.json()
                    content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
                    if content and len(content.strip()) > 10:
                        reply_text = content.strip()
                        used_model = model_name
                        break
            except Exception as e:
                print(f"[Chat] Model {model_name} failed: {e}")
                continue

    if not reply_text:
        # Grounded fallback if OpenRouter is unreachable or rate limited
        last_user_msg = messages[-1].get("content", "") if messages else ""
        reply_text = (
            f"Based on the technical findings (Verdict: {label}, Confidence: {round(confidence * 100)}%), "
            f"the ensemble detected {artifacts_summary}. Regarding your note ('{last_user_msg[:120]}'), "
            f"this context aligns with our detector signals: external compression or generator specifics often modulate "
            f"the high-frequency pixel distribution and EXIF integrity. Continued manual review is recommended."
        )
        used_model = "deepguard-offline-forensics"

    return {
        "reply": reply_text,
        "model": used_model,
    }