import os
import sys
import secrets
import spaces
import torch
import gradio as gr
from typing import List, Union, Optional, Dict, Any
from fastapi import APIRouter, Header, HTTPException, status
from pydantic import BaseModel
from huggingface_hub import snapshot_download
from sentence_transformers import SentenceTransformer

EMBED_MODEL_ID = "Dibachain/Diba-Embed"
MIZAN_MODEL_ID = "Nimruz/Mizan-v3-3B"
API_KEY = os.environ.get("MY_API_KEY", "")
device = "cuda" if torch.cuda.is_available() else "cpu"

# 1. Load Diba-Embed and register with ZeroGPU
print(f"Loading {EMBED_MODEL_ID}...")
embed_model = SentenceTransformer(
    EMBED_MODEL_ID,
    trust_remote_code=True,
    model_kwargs={"torch_dtype": torch.bfloat16}
)
embed_model.max_seq_length = 512
if torch.cuda.is_available():
    embed_model = embed_model.to("cuda")
print("Diba-Embed loaded successfully!")

# 2. Download and Load Nimruz/Mizan-v3-3B
print(f"Downloading and loading {MIZAN_MODEL_ID}...")
mizan_path = snapshot_download(MIZAN_MODEL_ID)
if mizan_path not in sys.path:
    sys.path.insert(0, mizan_path)

from joint_schema_model import load_release_model, systemone

target_device = "cuda" if torch.cuda.is_available() else "cpu"
mizan_model, mizan_processor = load_release_model(mizan_path, device=target_device)
print("Mizan-v3-3B loaded successfully!")

# 3. API Key Helper
def verify_key(provided_key: Optional[str]):
    if not API_KEY:
        raise RuntimeError("MY_API_KEY secret is not configured in Space Settings.")
    if not provided_key or not secrets.compare_digest(provided_key, API_KEY):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or missing API Key"
        )

# 4. ZeroGPU Function 1: Embeddings
@spaces.GPU(duration=30)
def compute_embeddings(
    api_key: str,
    texts: Union[str, List[str]],
    input_type: str = "document",
    normalize: bool = True
):
    if not API_KEY or not secrets.compare_digest(api_key or "", API_KEY):
        raise gr.Error("Unauthorized: Invalid API Key")

    text_list = [texts] if isinstance(texts, str) else texts
    prompt_name = "query" if input_type == "query" else None

    embeddings = embed_model.encode(
        text_list,
        prompt_name=prompt_name,
        normalize_embeddings=normalize,
        batch_size=8
    )
    return [
        {"object": "embedding", "index": i, "embedding": emb.tolist()}
        for i, emb in enumerate(embeddings)
    ]

# 5. ZeroGPU Function 2: Mizan-v3-3B Decision Inference
@spaces.GPU(duration=30)
def compute_decision(
    api_key: str,
    state: Any,
    questions: Dict[str, Any]
):
    if not API_KEY or not secrets.compare_digest(api_key or "", API_KEY):
        raise gr.Error("Unauthorized: Invalid API Key")

    payload = {
        "model": MIZAN_MODEL_ID,
        "state": state,
        "questions": questions
    }
    with torch.no_grad():
        return systemone(mizan_model, mizan_processor, payload)

# 6. Build a Tabbed Gradio Interface so ZeroGPU scans and registers BOTH @spaces.GPU functions
embed_ui = gr.Interface(
    fn=compute_embeddings,
    inputs=[
        gr.Textbox(label="API Key", type="password"),
        gr.Textbox(label="Input Text"),
        gr.Radio(["document", "query"], value="document", label="Input Type"),
        gr.Checkbox(value=True, label="Normalize")
    ],
    outputs=gr.JSON(label="Embeddings Output"),
    api_name=False
)

mizan_ui = gr.Interface(
    fn=compute_decision,
    inputs=[
        gr.Textbox(label="API Key", type="password"),
        gr.Textbox(label="State (Text or JSON)"),
        gr.JSON(label="Questions Dict")
    ],
    outputs=gr.JSON(label="Decision Output"),
    api_name=False
)

demo = gr.TabbedInterface(
    [embed_ui, mizan_ui],
    ["Diba-Embed", "Mizan-v3-3B"],
    title="Diba-Embed & Mizan-v3-3B Private API"
)

# Launch Gradio with ssr_mode=False so FastAPI binds directly to port 7860
demo.launch(prevent_thread_lock=True, ssr_mode=False)

# 7. Create an APIRouter for custom endpoints
api_router = APIRouter()

class EmbedRequest(BaseModel):
    input: Union[str, List[str]]
    input_type: Optional[str] = "document"
    normalize: Optional[bool] = True

class DecisionRequest(BaseModel):
    model: Optional[str] = MIZAN_MODEL_ID
    state: Any
    questions: Dict[str, Any]

def extract_token(authorization: Optional[str], x_api_key: Optional[str]) -> Optional[str]:
    if x_api_key:
        return x_api_key
    if authorization and authorization.lower().startswith("bearer "):
        return authorization[7:].strip()
    return None

@api_router.get("/health")
def health_check():
    return {
        "status": "online",
        "models": [EMBED_MODEL_ID, MIZAN_MODEL_ID],
        "endpoints": ["/v1/embeddings", "/v1/systemone", "/v1/decide"]
    }

@api_router.post("/v1/embeddings")
def create_embeddings(
    req: EmbedRequest,
    authorization: Optional[str] = Header(default=None),
    x_api_key: Optional[str] = Header(default=None)
):
    token = extract_token(authorization, x_api_key)
    verify_key(token)

    data = compute_embeddings(token, req.input, req.input_type or "document", req.normalize)
    return {
        "object": "list",
        "model": EMBED_MODEL_ID,
        "data": data
    }

@api_router.post("/v1/systemone")
@api_router.post("/v1/decide")
def create_decision(
    req: DecisionRequest,
    authorization: Optional[str] = Header(default=None),
    x_api_key: Optional[str] = Header(default=None)
):
    token = extract_token(authorization, x_api_key)
    verify_key(token)

    return compute_decision(token, req.state, req.questions)

# Prepend custom routes BEFORE Gradio's catch-all GET /{path:path} route
demo.app.router.routes[:0] = api_router.routes

demo.block_thread()