import os, json
from typing import Any, Dict, List, Optional, Tuple
import httpx


OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")
OPENROUTER_MODEL = os.environ.get("OPENROUTER_MODEL", "mistralai/mistral-small-24b-instruct-2501")

def _extract_json_object(text: str) -> Optional[dict]:
    """
    Best-effort: find the first top-level JSON object in a response.
    Works even if the model wraps JSON with prose or code fences.
    """
    if not text:
        return None

    # Strip common fences
    t = text.strip()
    t = t.removeprefix("```json").removeprefix("```").removesuffix("```").strip()

    # Find first {...} span
    start = t.find("{")
    if start == -1:
        return None

    depth = 0
    for i in range(start, len(t)):
        if t[i] == "{":
            depth += 1
        elif t[i] == "}":
            depth -= 1
            if depth == 0:
                chunk = t[start:i+1]
                try:
                    return json.loads(chunk)
                except Exception:
                    return None
    return None



def openrouter_chat(
    messages: List[Dict[str, str]],
    response_format: Optional[Dict[str, Any]] = None,
    temperature: float = 0.2,
    max_tokens: int = 512,
    timeout_s: float = 30.0,
) -> Tuple[Optional[str], Optional[dict], Optional[str]]:
    """
    Returns (raw_text, parsed_json, error_str).
    Never raises.

    Instrumented to help diagnose moderation/routing variance:
    - includes HTTP status
    - includes OpenRouter error message/code if provided
    """
    if not OPENROUTER_API_KEY:
        return None, None, "OPENROUTER_API_KEY missing"

    payload: Dict[str, Any] = {
        "model": OPENROUTER_MODEL,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    if response_format is not None:
        payload["response_format"] = response_format

    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://huggingface.co/spaces",
        "X-Title": "Prompt_Squirrel_RAG",
    }

    try:
        with httpx.Client(timeout=timeout_s) as client:
            r = client.post(
                "https://openrouter.ai/api/v1/chat/completions",
                headers=headers,
                json=payload,
            )

            try:
                data = r.json()
            except Exception:
                data = None

            if r.status_code >= 400:
                error_parts = [f"HTTP {r.status_code}"]
                if isinstance(data, dict):
                    err_obj = data.get("error")
                    if isinstance(err_obj, dict):
                        code = err_obj.get("code")
                        message = err_obj.get("message")
                        if code:
                            error_parts.append(f"code={code}")
                        if message:
                            error_parts.append(str(message))
                if len(error_parts) == 1:
                    body = (r.text or "").strip()
                    if body:
                        error_parts.append(body[:300])
                return None, None, "OpenRouter error: " + " | ".join(error_parts)

            if not isinstance(data, dict):
                return None, None, f"OpenRouter error: HTTP {r.status_code} | non-JSON response"

            choices = data.get("choices")
            if not isinstance(choices, list) or not choices:
                return None, None, f"OpenRouter error: HTTP {r.status_code} | missing choices"

            choice0 = choices[0] if isinstance(choices[0], dict) else {}
            message = choice0.get("message") if isinstance(choice0, dict) else {}
            content = (message.get("content", "") if isinstance(message, dict) else "") or ""
            content = content.strip()

            finish_reason = choice0.get("finish_reason") if isinstance(choice0, dict) else None
            native_finish_reason = choice0.get("native_finish_reason") if isinstance(choice0, dict) else None

            # (optional) expose these as part of error_str for logging
            meta = []
            if data.get("model"):
                meta.append(f"model={data['model']}")
            if finish_reason:
                meta.append(f"finish={finish_reason}")
            if native_finish_reason:
                meta.append(f"native_finish={native_finish_reason}")
            if isinstance(data.get("usage"), dict):
                u = data["usage"]
                if "prompt_tokens" in u and "completion_tokens" in u:
                    meta.append(f"tokens={u['prompt_tokens']}+{u['completion_tokens']}")

            parsed = _extract_json_object(content)

            # If it looks filtered, flag it
            if finish_reason == "content_filter":
                return content, parsed, f"Filtered (content_filter; {'; '.join(meta)})"

            # If it looks refusal-like but not content_filter, still flag it
            if content.lower().startswith(("i can't", "i can’t", "i cannot", "can't", "cannot")):
                return content, parsed, f"Refusal-like ({'; '.join(meta)})"

            return content, parsed, None

    except Exception as e:
        return None, None, f"{type(e).__name__}: {e}"
        
        
if __name__ == "__main__":
    print("openrouter_client.py imports ok")