"""
Multi-model consensus engine for the CSC Engine research lab protocol.

Runs observer + builder through multiple LLM providers, compares outputs,
flags disagreements, and produces a consensus verdict with timestamps.

Architecture:
  1. Call N providers for the same prompt
  2. Extract structured sections (EVIDENCE, REASONING, CODE) from each
  3. Compare section similarity and content overlap
  4. Flag disagreements (different code approaches, conflicting evidence)
  5. Produce consensus verdict: AGREED, PARTIAL_AGREEMENT, DISAGREEMENT
  6. Select best output (longest code, most evidence items, or majority vote)
  7. Every step timestamped and logged
"""

import os
import time
import json
import hashlib
import difflib
from typing import Optional


def _get_available_providers() -> list[str]:
    """Determine which LLM providers are available based on env vars."""
    providers = []
    if os.getenv("GROQ_API_KEY") or os.getenv("GROK_API_KEY"):
        providers.append("groq")
    if os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN"):
        providers.append("huggingface")
    if os.getenv("OPENAI_API_KEY"):
        providers.append("openai")
    if os.getenv("OLLAMA_HOST"):
        providers.append("ollama")
    # Always include the primary provider
    primary = os.getenv("PROVIDER", "groq").lower().strip()
    if primary == "hybrid":
        primary = "groq"
    if primary not in providers:
        providers.insert(0, primary)
    # Deduplicate, preserve order, max 3
    seen = set()
    unique = []
    for p in providers:
        if p not in seen:
            seen.add(p)
            unique.append(p)
    return unique[:3]


def _call_provider(provider: str, prompt: str, frame_b64: Optional[str] = None) -> str:
    """Call a single LLM provider and return its output."""
    provider = provider.lower().strip()
    if provider == "groq" or provider == "grok":
        from .observer_llm import call_grok
        return call_grok(prompt, frame_b64)
    elif provider == "huggingface":
        from .observer_llm import call_hf_inference
        return call_hf_inference(prompt, frame_b64)
    elif provider == "openai":
        from .observer_llm import call_openai
        return call_openai(prompt, frame_b64)
    elif provider == "ollama":
        from .observer_llm import call_ollama
        return call_ollama(prompt, frame_b64)
    else:
        raise RuntimeError(f"Unknown provider: {provider}")


def _call_builder_provider(provider: str, prompt: str) -> str:
    """Call a single builder LLM provider."""
    provider = provider.lower().strip()
    if provider == "groq" or provider == "grok":
        from .builder_llm import call_grok
        return call_grok(prompt)
    elif provider == "huggingface":
        from .builder_llm import call_hf_inference
        return call_hf_inference(prompt)
    elif provider == "openai":
        from .builder_llm import call_openai
        return call_openai(prompt)
    elif provider == "ollama":
        from .builder_llm import call_ollama
        return call_ollama(prompt)
    else:
        raise RuntimeError(f"Unknown builder provider: {provider}")


def _extract_section(text: str, section_name: str) -> str:
    """Extract a named section (EVIDENCE, REASONING, CODE, etc.) from LLM output."""
    lines = text.split("\n")
    capturing = False
    collected = []
    for line in lines:
        stripped = line.strip().upper()
        if stripped.startswith(section_name + ":"):
            capturing = True
            continue
        if capturing:
            # Check if we hit another section header
            for header in ["EVIDENCE:", "REASONING:", "CODE:", "RUN:", "TEST:", "ATTRIBUTION:", "SCENE:", "INFERRED_INTENT:", "SIGNALS:", "CANDIDATE_TASK:", "UNCERTAINTY:"]:
                if stripped.startswith(header):
                    capturing = False
                    break
            if capturing:
                collected.append(line)
    return "\n".join(collected).strip()


def _similarity(text_a: str, text_b: str) -> float:
    """Compute text similarity ratio between two strings (0.0 to 1.0)."""
    if not text_a or not text_b:
        return 0.0
    return difflib.SequenceMatcher(None, text_a.lower(), text_b.lower()).ratio()


def _code_similarity(code_a: str, code_b: str) -> float:
    """Compare code similarity ignoring whitespace and comments."""
    def normalize(code: str) -> str:
        lines = []
        for line in code.split("\n"):
            line = line.strip()
            if line and not line.startswith("#"):
                lines.append(line)
        return " ".join(lines)
    return _similarity(normalize(code_a), normalize(code_b))


def observer_consensus(state, compact_state: dict) -> dict:
    """Run observer through multiple providers and produce consensus.

    Returns dict with:
    - consensus_verdict: AGREED | PARTIAL_AGREEMENT | DISAGREEMENT | SINGLE_PROVIDER
    - observer_output: best output selected
    - all_outputs: list of {provider, output, timestamp}
    - disagreements: list of flagged differences
    - similarity_matrix: pairwise similarities
    - timestamp: consensus timestamp
    """
    from .observer_llm import OBSERVER_PROMPT
    from .stream_state import SessionState

    frames = list(state.frames)
    last_frame = frames[-1] if frames else None
    frame_b64 = last_frame.jpeg_b64 if last_frame else None

    prompt = f"""{OBSERVER_PROMPT}

COMPRESSED SENSORY STATE:
{json.dumps(compact_state, indent=2)}

Produce your observation now."""

    providers = _get_available_providers()
    timestamp_start = time.time()

    all_outputs = []
    errors = []

    for provider in providers:
        ts = time.time()
        try:
            output = _call_provider(provider, prompt, frame_b64)
            all_outputs.append({
                "provider": provider,
                "output": output,
                "timestamp": ts,
                "duration_ms": int((time.time() - ts) * 1000),
                "error": None,
            })
        except Exception as e:
            errors.append({
                "provider": provider,
                "error": str(e)[:200],
                "timestamp": ts,
            })

    if not all_outputs:
        raise RuntimeError(f"All observer providers failed: {errors}")

    # If only one provider succeeded, return single-provider verdict
    if len(all_outputs) == 1:
        return {
            "consensus_verdict": "SINGLE_PROVIDER",
            "observer_output": all_outputs[0]["output"],
            "all_outputs": all_outputs,
            "errors": errors,
            "disagreements": [],
            "similarity_matrix": {},
            "providers_used": [all_outputs[0]["provider"]],
            "timestamp": time.time(),
            "duration_ms": int((time.time() - timestamp_start) * 1000),
        }

    # Compare outputs pairwise
    n = len(all_outputs)
    similarity_matrix = {}
    disagreements = []

    for i in range(n):
        for j in range(i + 1, n):
            out_a = all_outputs[i]["output"]
            out_b = all_outputs[j]["output"]
            prov_a = all_outputs[i]["provider"]
            prov_b = all_outputs[j]["provider"]

            # Compare overall similarity
            overall_sim = _similarity(out_a, out_b)

            # Compare specific sections
            scene_a = _extract_section(out_a, "SCENE")
            scene_b = _extract_section(out_b, "SCENE")
            intent_a = _extract_section(out_a, "INFERRED_INTENT")
            intent_b = _extract_section(out_b, "INFERRED_INTENT")

            scene_sim = _similarity(scene_a, scene_b)
            intent_sim = _similarity(intent_a, intent_b)

            key = f"{prov_a}_vs_{prov_b}"
            similarity_matrix[key] = {
                "overall": round(overall_sim, 3),
                "scene": round(scene_sim, 3),
                "intent": round(intent_sim, 3),
            }

            # Flag disagreements
            if overall_sim < 0.3:
                disagreements.append({
                    "providers": [prov_a, prov_b],
                    "type": "low_overall_similarity",
                    "similarity": round(overall_sim, 3),
                    "detail": "Outputs differ significantly in content and structure",
                })
            if intent_sim < 0.4 and intent_a and intent_b:
                disagreements.append({
                    "providers": [prov_a, prov_b],
                    "type": "intent_divergence",
                    "similarity": round(intent_sim, 3),
                    "detail_a": intent_a[:200],
                    "detail_b": intent_b[:200],
                })

    # Determine consensus verdict
    avg_sim = sum(s["overall"] for s in similarity_matrix.values()) / len(similarity_matrix) if similarity_matrix else 0
    if avg_sim >= 0.6:
        verdict = "AGREED"
    elif avg_sim >= 0.3:
        verdict = "PARTIAL_AGREEMENT"
    else:
        verdict = "DISAGREEMENT"

    # Select best output: prefer the one with most content (longest output)
    best = max(all_outputs, key=lambda x: len(x["output"]))

    return {
        "consensus_verdict": verdict,
        "observer_output": best["output"],
        "all_outputs": all_outputs,
        "errors": errors,
        "disagreements": disagreements,
        "similarity_matrix": similarity_matrix,
        "avg_similarity": round(avg_sim, 3),
        "providers_used": [o["provider"] for o in all_outputs],
        "timestamp": time.time(),
        "duration_ms": int((time.time() - timestamp_start) * 1000),
    }


def builder_consensus(state, observation: dict, gate_result=None) -> dict:
    """Run builder through multiple providers and produce consensus.

    Returns dict with:
    - consensus_verdict: AGREED | PARTIAL_AGREEMENT | DISAGREEMENT | SINGLE_PROVIDER
    - patch_output: best output selected
    - patch_hash: hash of selected output
    - receipt: receipt with consensus metadata
    - all_outputs: list of {provider, output, code_extracted, timestamp}
    - disagreements: list of flagged code differences
    - similarity_matrix: pairwise code similarities
    """
    from .builder_llm import BUILDER_PROMPT, _extract_reasons, _extract_uncertainty
    from .stream_state import sha256_text

    fallback_level = 1
    artifact_type = "task_code"
    sensory_channels = []
    feature_attribution = {}
    if gate_result:
        fallback_level = gate_result.fallback_level
        artifact_type = gate_result.artifact_type
        sensory_channels = gate_result.sensory_channels
        feature_attribution = gate_result.feature_attribution

    prompt = f"""{BUILDER_PROMPT}

OBSERVER OUTPUT:
{observation.get('observer_output', '')}

MODE: {state.mode}

FALLBACK LEVEL: {fallback_level}
ARTIFACT TYPE: {artifact_type}
SENSORY CHANNELS: {', '.join(sensory_channels) if sensory_channels else 'none'}
FEATURE ATTRIBUTION: {json.dumps(feature_attribution, default=str)}

You are at FALLBACK LEVEL {fallback_level}. Generate {artifact_type} based on the available sensory evidence. NEVER return INSUFFICIENT_EVIDENCE. Always produce a CODE section with runnable Python."""

    providers = _get_available_providers()
    timestamp_start = time.time()

    all_outputs = []
    errors = []

    for provider in providers:
        ts = time.time()
        try:
            output = _call_builder_provider(provider, prompt)
            code = _extract_section(output, "CODE")
            # Strip markdown code fences
            if code.startswith("```"):
                code = "\n".join(code.split("\n")[1:])
            if code.endswith("```"):
                code = code.rsplit("```", 1)[0]
            all_outputs.append({
                "provider": provider,
                "output": output,
                "code_extracted": code.strip(),
                "code_lines": len([l for l in code.strip().split("\n") if l.strip()]),
                "timestamp": ts,
                "duration_ms": int((time.time() - ts) * 1000),
                "error": None,
            })
        except Exception as e:
            errors.append({
                "provider": provider,
                "error": str(e)[:200],
                "timestamp": ts,
            })

    if not all_outputs:
        raise RuntimeError(f"All builder providers failed: {errors}")

    if len(all_outputs) == 1:
        best = all_outputs[0]
        patch_hash = sha256_text(best["output"] + str(time.time()))
        receipt = {
            "receipt_type": "PATCH_RECEIPT_V1",
            "patch_hash": patch_hash,
            "session_id": state.session_id,
            "timestamp": time.time(),
            "derived_from": {
                "frame_hashes": state.frame_hashes(),
                "audio_chunk_hashes": state.audio_chunk_hashes(),
                "speaker_segments": state.speaker_segments(),
                "observer_state_hash": observation.get("state_hash", ""),
            },
            "reason_codes": _extract_reasons(state, observation),
            "uncertainty": _extract_uncertainty(observation),
            "mode": state.mode,
            "provider": best["provider"],
            "fallback_level": fallback_level,
            "artifact_type": artifact_type,
            "sensory_channels": sensory_channels,
            "feature_attribution": feature_attribution,
            "consensus": {
                "verdict": "SINGLE_PROVIDER",
                "providers_used": [best["provider"]],
            },
        }
        return {
            "consensus_verdict": "SINGLE_PROVIDER",
            "patch_output": best["output"],
            "patch_hash": patch_hash,
            "receipt": receipt,
            "all_outputs": all_outputs,
            "errors": errors,
            "disagreements": [],
            "similarity_matrix": {},
            "fallback_level": fallback_level,
            "artifact_type": artifact_type,
        }

    # Compare code outputs pairwise
    n = len(all_outputs)
    similarity_matrix = {}
    disagreements = []

    for i in range(n):
        for j in range(i + 1, n):
            code_a = all_outputs[i]["code_extracted"]
            code_b = all_outputs[j]["code_extracted"]
            prov_a = all_outputs[i]["provider"]
            prov_b = all_outputs[j]["provider"]

            code_sim = _code_similarity(code_a, code_b)

            key = f"{prov_a}_vs_{prov_b}"
            similarity_matrix[key] = {
                "code_similarity": round(code_sim, 3),
                "lines_a": all_outputs[i]["code_lines"],
                "lines_b": all_outputs[j]["code_lines"],
            }

            if code_sim < 0.3:
                disagreements.append({
                    "providers": [prov_a, prov_b],
                    "type": "different_code_approaches",
                    "similarity": round(code_sim, 3),
                    "detail": "Providers produced substantially different code implementations",
                })
            elif code_sim < 0.6:
                disagreements.append({
                    "providers": [prov_a, prov_b],
                    "type": "partial_code_divergence",
                    "similarity": round(code_sim, 3),
                    "detail": "Providers produced similar but not identical code",
                })

    avg_sim = sum(s["code_similarity"] for s in similarity_matrix.values()) / len(similarity_matrix) if similarity_matrix else 0
    if avg_sim >= 0.6:
        verdict = "AGREED"
    elif avg_sim >= 0.3:
        verdict = "PARTIAL_AGREEMENT"
    else:
        verdict = "DISAGREEMENT"

    # Select best output: prefer longest code (most complete implementation)
    best = max(all_outputs, key=lambda x: x["code_lines"])

    patch_hash = sha256_text(best["output"] + str(time.time()))
    receipt = {
        "receipt_type": "PATCH_RECEIPT_V1",
        "patch_hash": patch_hash,
        "session_id": state.session_id,
        "timestamp": time.time(),
        "derived_from": {
            "frame_hashes": state.frame_hashes(),
            "audio_chunk_hashes": state.audio_chunk_hashes(),
            "speaker_segments": state.speaker_segments(),
            "observer_state_hash": observation.get("state_hash", ""),
        },
        "reason_codes": _extract_reasons(state, observation),
        "uncertainty": _extract_uncertainty(observation),
        "mode": state.mode,
        "provider": best["provider"],
        "fallback_level": fallback_level,
        "artifact_type": artifact_type,
        "sensory_channels": sensory_channels,
        "feature_attribution": feature_attribution,
        "consensus": {
            "verdict": verdict,
            "providers_used": [o["provider"] for o in all_outputs],
            "avg_code_similarity": round(avg_sim, 3),
            "disagreements_count": len(disagreements),
            "selected_provider": best["provider"],
            "selection_criteria": "most_code_lines",
        },
    }

    return {
        "consensus_verdict": verdict,
        "patch_output": best["output"],
        "patch_hash": patch_hash,
        "receipt": receipt,
        "all_outputs": all_outputs,
        "errors": errors,
        "disagreements": disagreements,
        "similarity_matrix": similarity_matrix,
        "avg_similarity": round(avg_sim, 3),
        "providers_used": [o["provider"] for o in all_outputs],
        "fallback_level": fallback_level,
        "artifact_type": artifact_type,
        "timestamp": time.time(),
        "duration_ms": int((time.time() - timestamp_start) * 1000),
    }