"""Model interface for calling moderation models."""

import json
import math
import re

from openai import OpenAI
from openai_harmony import (
    DeveloperContent,
    HarmonyEncodingName,
    Message,
    Role,
    SystemContent,
    load_harmony_encoding,
)

from utils.constants import (
    DEFAULT_MAX_TOKENS,
    DEFAULT_TEMPERATURE,
    DEFAULT_TOP_P,
    GENERIC_SYSTEM_PROMPT_PREFIX,
    LM_PROMPT_INSTRUCT,
    COPE_PROMPT_TEMPLATE,
    RESPONSE_FORMAT,
    ROUTER_URL,
    MODELS,
    COPE_B_MODEL_ID,
    COPE_B_THRESHOLD,
    MILA_MODEL_ID,
    MILA_THRESHOLD,
    SHIELDSTRAL_MODEL_ID,
    SHIELDSTRAL_SYSTEM_PROMPT,
    SHIELDSTRAL_USER_TEMPLATE,
    SHIELDSTRAL_THRESHOLD,
    POLICYLM_MODEL_ID,
    POLICYLM_CUTOFFS,
    POLICYLM_DEFAULT_PRESET,
)


def get_model_info(model_id: str) -> dict:
    """Get model metadata by ID."""
    for model in MODELS:
        if model["id"] == model_id:
            return model
    return None


def extract_model_id(choice: str) -> str:
    """Extract model ID from dropdown choice format 'Name (id)'."""
    if not choice:
        return ""
    return choice.split("(")[-1].rstrip(")")


def is_gptoss_model(model_id: str) -> bool:
    """Check if model is GPT-OSS."""
    return model_id.startswith("openai/gpt-oss")


def get_default_system_prompt(model_id: str, reasoning_effort: str = "Low") -> str:
    """Generate default system prompt based on model type and policy."""
    if is_gptoss_model(model_id):
        enc = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
        system_prompt_harmony = Message.from_role_and_content(
            Role.SYSTEM, SystemContent.new().with_reasoning_effort(reasoning_effort)
        )
        system_prompt_dict = enc.decode(enc.render(system_prompt_harmony))
        system_prompt_content = re.search(r"<\|message\|>(.*?)<\|end\|>", system_prompt_dict, re.DOTALL).group(1)
        return system_prompt_content
    else:
        # PolicyLM, Shieldstral, and CoPE build their own fixed prompts (see each
        # model card) and never receive this field, so there is no default to offer.
        return ""


def make_messages(test: str, policy: str, model_id: str, reasoning_effort: str = "Low", system_prompt: str | None = None, response_format: str = RESPONSE_FORMAT) -> list[dict]:
    """Create messages based on model type."""
    if model_id.startswith("openai/gpt-oss-safeguard"):
        # GPT-OSS uses Harmony encoding
        enc = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
        system_content = SystemContent.new().with_reasoning_effort(reasoning_effort)
        # The System Prompt box is editable, and this model takes whatever is in it as
        # its system message verbatim. Cleared, that sent content: null and the router
        # rejected the whole request with "Value is not nullable" — an error naming
        # nothing the user could connect to the empty box they had just emptied. Fall
        # back to the harmony-rendered system message, which is what the box is
        # pre-filled with anyway.
        if not (system_prompt or "").strip():
            rendered = enc.decode(
                enc.render(Message.from_role_and_content(Role.SYSTEM, system_content))
            )
            system_prompt = re.search(r"<\|message\|>(.*?)<\|end\|>", rendered, re.DOTALL).group(1)
        conv_messages = [
            Message.from_role_and_content(
                Role.DEVELOPER,
                DeveloperContent.new().with_instructions(policy + "\n\n" + response_format),
            ),
            Message.from_role_and_content(Role.USER, test),
        ]
        messages = [
            {"role": "system", "content": system_prompt},
        ]
        for pre_msg in conv_messages:
            tokens = enc.render(pre_msg)
            prompt = enc.decode(tokens)
            messages.append({
                "role": re.search(r"<\|start\|>(.*?)<\|message\|>", prompt).group(1),
                "content": re.search(r"<\|message\|>(.*?)<\|end\|>", prompt, re.DOTALL).group(1),
            })
        return messages
    else:
        system_content = LM_PROMPT_INSTRUCT.format(
            system_prompt=system_prompt,
            policy=policy,
            response_format=response_format
        )
        return [
            {"role": "system", "content": system_content},
            {"role": "user", "content": f"Content: {test}\n\nResponse:"},
        ]


def _normalize_token(tok: str) -> str:
    """Lowercase and strip whitespace / leading sentinel chars from a token."""
    return tok.strip().lstrip("Ġ▁ ").lower()


def run_shieldstral(test_input: str, policy: str) -> dict:
    """
    Run Shieldstral-1.0-3B against its dedicated HF Inference Endpoint.

    Shieldstral emits a single yes/no token; we softmax the yes vs. no logprobs
    into a 0-1 violation probability (yes == violation) and threshold at 0.5.
    """
    import gradio as gr

    from utils.helpers import get_shieldstral_config

    endpoint_url, token, served_name = get_shieldstral_config()
    if not endpoint_url or not token:
        raise gr.Error(
            "Shieldstral endpoint not configured. Set the SHIELDSTRAL_ENDPOINT_URL and "
            "SHIELDSTRAL_INFERENCE_TOKEN Space secrets (see README)."
        )

    client = OpenAI(base_url=endpoint_url, api_key=token)
    messages = [
        {"role": "system", "content": SHIELDSTRAL_SYSTEM_PROMPT},
        {"role": "user", "content": SHIELDSTRAL_USER_TEMPLATE.format(policy=policy, content=test_input)},
    ]

    try:
        completion = client.chat.completions.create(
            model=served_name,
            messages=messages,
            max_tokens=1,
            temperature=0.0,
            logprobs=True,
            top_logprobs=20,
        )
    except Exception as e:
        raise gr.Error(f"Shieldstral inference failed: {e}. Check the endpoint is running and the token is valid.")

    try:
        top = completion.choices[0].logprobs.content[0].top_logprobs
    except (AttributeError, IndexError, TypeError):
        raise gr.Error(
            "Shieldstral endpoint did not return token logprobs. Ensure the endpoint is served with "
            "logprobs enabled (vLLM OpenAI-compatible server)."
        )

    z_yes = max((lp.logprob for lp in top if _normalize_token(lp.token).startswith("yes")), default=None)
    z_no = max((lp.logprob for lp in top if _normalize_token(lp.token).startswith("no")), default=None)

    if z_yes is None and z_no is None:
        raise gr.Error("Shieldstral returned neither a 'yes' nor 'no' token in the top logprobs.")
    if z_yes is None:
        score = 0.0
    elif z_no is None:
        score = 1.0
    else:
        score = math.exp(z_yes) / (math.exp(z_yes) + math.exp(z_no))

    label = 1 if score > SHIELDSTRAL_THRESHOLD else 0
    # Shieldstral is a single-score model with no taxonomy, but an empty list makes the
    # UI report "No categories found in response" and suggest the model returned bad
    # JSON. Report the one decision it does make.
    result_content = json.dumps({
        "label": label,
        "categories": [{
            "category": "Policy violation",
            "reasoning": f"violation probability {score:.4f} (threshold {SHIELDSTRAL_THRESHOLD})",
        }],
    })
    reasoning = (
        f"Shieldstral violation probability: {score:.4f} "
        f"(threshold {SHIELDSTRAL_THRESHOLD}, yes==violation).\n"
        f"Raw top token: {completion.choices[0].message.content!r}"
    )
    return {"content": f"```json\n{result_content}\n```", "reasoning": reasoning}


def run_policylm(test_input: str, policy: str, structured_rules: list | None = None) -> dict:
    """
    Run PolicyLM-1.7B against its dedicated HF Inference Endpoint.

    PolicyLM is not OpenAI-compatible: the endpoint must run the custom handler in
    policylm_handler.py, which wraps the model's helper library and returns
    {"flagged": bool, "max_score": float, "categories": [{"category","score","flagged"}]}.

    When `structured_rules` are supplied (the "flag / don't flag / exception" rules
    from the Policy tab), they are sent as discrete `categories` so the handler builds
    one PolicyLM Category per rule — exceptions are read as exceptions, and each rule
    gets its own score. Without them, the handler falls back to wrapping `policy` as a
    single catch-all category (legacy behaviour).
    """
    import gradio as gr
    import requests

    from utils.helpers import get_policylm_config
    from utils.policy_rules import policy_text_to_categories, rules_to_categories

    endpoint_url, token = get_policylm_config()
    if not endpoint_url or not token:
        raise gr.Error(
            "PolicyLM endpoint not configured. Set the POLICYLM_ENDPOINT_URL and "
            "POLICYLM_INFERENCE_TOKEN Space secrets, and deploy the endpoint with policylm_handler.py (see README)."
        )

    inputs = {"message": test_input, "policy": policy, "preset": POLICYLM_DEFAULT_PRESET}
    # Always send categories. The handler's fallback wraps the whole policy document in
    # one category, and PolicyLM rejects any field containing a newline
    # ("Violation rule contains a line break; each field is one line"), so a multi-line
    # written policy fails there. Parsing it here also recovers per-rule scores.
    categories_payload = rules_to_categories(structured_rules) or policy_text_to_categories(policy)
    if categories_payload:
        inputs["categories"] = categories_payload

    try:
        resp = requests.post(
            endpoint_url,
            headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
            json={"inputs": inputs},
            timeout=120,
        )
    except Exception as e:
        raise gr.Error(f"PolicyLM request failed: {e}. Check the endpoint is running.")

    if resp.status_code != 200:
        raise gr.Error(f"PolicyLM endpoint returned HTTP {resp.status_code}: {resp.text[:500]}")

    try:
        data = resp.json()
    except Exception:
        raise gr.Error(f"PolicyLM endpoint returned non-JSON response: {resp.text[:500]}")

    # Our handler always returns ≥1 category; anything else must not read as "no violation".
    if not isinstance(data, dict) or "flagged" not in data or not data.get("categories"):
        raise gr.Error(
            "PolicyLM endpoint returned an unexpected response (expected flagged + non-empty "
            "categories; is it running policylm_handler.py as handler.py?): "
            f"{resp.text[:500]}"
        )

    flagged = bool(data["flagged"])
    categories = data["categories"]
    label = 1 if flagged else 0

    cutoff = POLICYLM_CUTOFFS.get(POLICYLM_DEFAULT_PRESET)
    result_content = json.dumps({
        "label": label,
        "categories": [
            {
                "category": c.get("category", "Unknown"),
                "reasoning": f"score {c.get('score', float('nan')):.4f}"
                             + (" (flagged)" if c.get("flagged") else ""),
            }
            for c in categories
        ],
    })
    max_score = data.get("max_score")
    reasoning_lines = [f"PolicyLM preset: {POLICYLM_DEFAULT_PRESET} (cutoff {cutoff})."]
    if max_score is not None:
        reasoning_lines.append(f"Max category score: {max_score:.4f}.")
    for c in categories:
        reasoning_lines.append(f"- {c.get('category', 'Unknown')}: {c.get('score', float('nan')):.4f}")
    return {"content": f"```json\n{result_content}\n```", "reasoning": "\n".join(reasoning_lines)}


def run_cope_b(test_input: str, policy: str) -> dict:
    """
    Run CoPE-B-S1 against its dedicated SGLang "decision model" HF Inference Endpoint.

    CoPE-B-S1 is NOT a standard generate() LLM — it implements the System One decision
    spec. The endpoint runs Zentropi's `zentropiai/cope-s1-sglang` container, which
    exposes POST /v1/systemone. We send the content as `state.content` plus a single
    `noul` (true/false) question wrapping the policy, following the model card's
    content-policy convention. The server reads the decision from a pooled prefill slot
    (no tokens generated) and returns P(true) = probability the content merits a label.
    """
    import gradio as gr
    import requests

    from utils.helpers import get_cope_b_config

    endpoint_url, token, _served_name = get_cope_b_config()
    if not endpoint_url or not token:
        raise gr.Error(
            "CoPE-B endpoint not configured. Set the COPE_B_ENDPOINT_URL and "
            "COPE_B_INFERENCE_TOKEN Space secrets (see README)."
        )

    url = endpoint_url.rstrip("/") + "/v1/systemone"
    body = {
        "questions": {
            "q1": {
                "type": "noul",
                "instructions": "Does the content meet the criteria described in the following policy?\n\n" + policy,
                "criteria": {
                    "true": "The content meets the policy's criteria.",
                    "false": "The content does not meet the policy's criteria.",
                },
            }
        },
        "state": {"content": test_input},
    }

    try:
        resp = requests.post(
            url,
            headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
            json=body,
            timeout=120,
        )
    except Exception as e:
        raise gr.Error(f"CoPE-B-S1 request failed: {e}. Check the endpoint is running.")

    if resp.status_code != 200:
        raise gr.Error(f"CoPE-B-S1 endpoint returned HTTP {resp.status_code}: {resp.text[:500]}")

    try:
        data = resp.json()
        p_violation = float(data["answers"]["q1"]["noul"])
    except Exception:
        raise gr.Error(
            f"CoPE-B-S1 returned an unexpected response (expected answers.q1.noul): {resp.text[:400]}"
        )

    label = 1 if p_violation >= COPE_B_THRESHOLD else 0
    # Same as Shieldstral: one decision, no taxonomy. An empty list reads to the UI as
    # a malformed response rather than as a model that simply has no categories.
    result_content = json.dumps({
        "label": label,
        "categories": [{
            "category": "Policy violation",
            "reasoning": f"P(violation) {p_violation:.4f} (threshold {COPE_B_THRESHOLD})",
        }],
    })
    reasoning = (
        f"CoPE-B-S1 P(violation) = {p_violation:.4f} "
        f"(threshold {COPE_B_THRESHOLD}; System One `noul` read, pooled prefill — no tokens generated)."
    )
    return {"content": f"```json\n{result_content}\n```", "reasoning": reasoning}


def run_mila(test_input: str) -> dict:
    """
    Run the Mila Suicide Prevention Output Guardrail against its dedicated endpoint.

    Takes no policy, by design: this model scores against its own fixed policy (the
    MLCommons AILuminate "Suicide and Self-Harm" category) and has no way to be
    steered by the room's. It is also commission-based — it flags text that *enables
    or encourages* self-harm, not text expressing distress — so on a policy that asks
    for distress to be routed to support it will legitimately disagree with the other
    models rather than be wrong.
    """
    import gradio as gr
    import requests

    from utils.helpers import get_mila_config

    endpoint_url, token = get_mila_config()
    if not endpoint_url or not token:
        raise gr.Error(
            "Mila endpoint not configured. Set the MILA_ENDPOINT_URL and "
            "MILA_INFERENCE_TOKEN Space secrets (see README)."
        )

    try:
        resp = requests.post(
            endpoint_url.rstrip("/"),
            headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
            json={"inputs": test_input, "parameters": {"threshold": MILA_THRESHOLD}},
            timeout=120,
        )
        resp.raise_for_status()
        payload = resp.json()
    except Exception as e:
        raise gr.Error(f"Mila inference failed: {e}. Check the endpoint is running and the token is valid.")

    if isinstance(payload, list):
        payload = payload[0] if payload else {}
    score = payload.get("p_unsafe", payload.get("score"))
    if not isinstance(score, (int, float)):
        raise gr.Error(
            f"Mila endpoint returned an unexpected response (expected a p_unsafe score): {payload}"
        )

    label = 1 if score >= MILA_THRESHOLD else 0
    result_content = json.dumps({
        "label": label,
        "categories": [{
            "category": "Enables or encourages suicide or self-harm",
            "reasoning": f"p(unsafe) {score:.4f} (threshold {MILA_THRESHOLD})",
        }],
    })
    reasoning = (
        f"Mila p(unsafe): {score:.4f} (threshold {MILA_THRESHOLD}).\n"
        "Scored against the model's own fixed self-harm policy (MLCommons AILuminate), "
        "NOT the policy loaded in this app — that text is not sent to this model.\n"
        "It flags only content that enables, encourages or endorses self-harm; expressions "
        "of distress and help-seeking are 'safe' by its policy even when your policy would "
        "route them to support."
    )
    return {"content": f"```json\n{result_content}\n```", "reasoning": reasoning}


def run_test(
    model_id: str,
    test_input: str,
    policy: str,
    hf_token: str,
    reasoning_effort: str = "Low",
    max_tokens: int = DEFAULT_MAX_TOKENS,
    temperature: float = DEFAULT_TEMPERATURE,
    top_p: float = DEFAULT_TOP_P,
    system_prompt: str | None = None,
    response_format: str = RESPONSE_FORMAT,
    structured_rules: list | None = None,
) -> dict:
    """Run test on model."""
    import gradio as gr
    
    model_info = get_model_info(model_id)
    if not model_info:
        raise gr.Error(f"Unknown model: {model_id}. Please select a valid model from the dropdown.")

    client = OpenAI(base_url=ROUTER_URL, api_key=hf_token)

    # CoPE special case as a completions-only model
    if model_id.startswith("zentropi-ai/cope-a-9b"):
        try:
            from utils.helpers import get_cope_token
            cope_token = get_cope_token()
            client = OpenAI(base_url='https://q0c7gn9b2s2wkb6v.us-east-1.aws.endpoints.huggingface.cloud/v1', api_key=cope_token)
            cope_prompt = COPE_PROMPT_TEMPLATE.format(policy=policy,content=test_input)
            response = client.completions.create(
                model='cope-a-9b',
                prompt=cope_prompt,
                max_tokens=1,
                temperature=0,
            )
            result_content = f"""```json
{{
  "label": {response.choices[0].text},
  "categories": []
}}
"""
            result = {"content": result_content}
            return result
        except Exception as e:
            error_msg = str(e)
            raise gr.Error(f"Model inference failed: {error_msg}. Sorry, it was a last minute hack to add this model to this interface!")

    # Dedicated-endpoint classifiers (not on the router): each has its own prompt
    # format and native output, mapped to the {label, categories} contract.
    if model_id == SHIELDSTRAL_MODEL_ID:
        return run_shieldstral(test_input, policy)
    if model_id == POLICYLM_MODEL_ID:
        return run_policylm(test_input, policy, structured_rules)
    if model_id == COPE_B_MODEL_ID:
        return run_cope_b(test_input, policy)
    if model_id == MILA_MODEL_ID:
        # No policy argument: this model cannot be steered by one.
        return run_mila(test_input)

    messages = make_messages(test_input, policy, model_id, reasoning_effort, system_prompt, response_format)

    try:
        completion = client.chat.completions.create(
            model=model_id,
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            stop=None,
            extra_headers={"X-HF-Bill-To": "roosttools"},
        )
    except Exception as e:
        error_msg = str(e)
        if "401" in error_msg or "authentication" in error_msg.lower():
            raise gr.Error(f"Authentication failed: {error_msg}. Please check your token permissions.")
        elif "400" in error_msg or "bad request" in error_msg.lower():
            raise gr.Error(f"Invalid request: {error_msg}. Please check your input and try again.")
        elif "429" in error_msg or "rate limit" in error_msg.lower():
            raise gr.Error(f"Rate limit exceeded: {error_msg}. Please wait a moment and try again.")
        elif "timeout" in error_msg.lower():
            raise gr.Error(f"Request timed out: {error_msg}. Please try again with a shorter input or lower max_tokens.")
        else:
            raise gr.Error(f"Model inference failed: {error_msg}. Please check your inputs and try again.")

    result = {"content": completion.choices[0].message.content}

    # Extract reasoning if available
    message = completion.choices[0].message
    if model_info["is_thinking"]:
        if is_gptoss_model(model_id):
            # GPT-OSS: check reasoning or reasoning_content field
            if hasattr(message, "reasoning") and message.reasoning:
                result["reasoning"] = message.reasoning
            elif hasattr(message, "reasoning_content") and message.reasoning_content:
                result["reasoning"] = message.reasoning_content
        else:
            # Qwen Thinking: extract from content using </think> tag
            content = message.content
            if "</think>" in content:
                result["reasoning"] = content.split("</think>")[0].strip()
                # Also update content to be the part after </think>
                result["content"] = content.split("</think>")[-1].strip()

    return result