import os
import sys
import warnings
warnings.filterwarnings("ignore")
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"

from pathlib import Path
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from contextlib import asynccontextmanager
from typing import Optional
import torch
import pickle
import re
import redis
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# ── Constants ────────────────────────────────────────────
CONFIDENCE_THRESHOLD = 10.0
OOS_INTENT           = "oos"
OOS_RESPONSE         = "I'm not sure I understand that. Could you rephrase or ask me something banking related?"
MAX_LEN              = 64
DEVICE               = torch.device("cpu")
MODEL_REPO           = "RexTheSlasher/intent-classifier-bert"

# ── Redis Cloud ──────────────────────────────────────────
r = redis.Redis(
    host=os.getenv('REDIS_HOST'),
    port=int(os.getenv('REDIS_PORT', 6379)),
    password=os.getenv('REDIS_PASSWORD'),
    decode_responses=True
)

# ── Load model from Hugging Face Hub ────────────────────
print("Loading model from Hugging Face...")
tokenizer     = AutoTokenizer.from_pretrained(MODEL_REPO)
model         = AutoModelForSequenceClassification.from_pretrained(MODEL_REPO)
model         = model.to(DEVICE)
model.eval()

# ── Load label encoder from Hugging Face Hub ────────────
encoder_path  = hf_hub_download(repo_id=MODEL_REPO, filename="label_encoder.pkl")
with open(encoder_path, "rb") as f:
    label_encoder = pickle.load(f)

print(f"Model loaded! Classes: {len(label_encoder.classes_)}")

# ── Text cleaning ────────────────────────────────────────
def clean_text(text: str) -> str:
    text = text.lower()
    text = re.sub(r"[^a-z0-9\s]", "", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text

# ── Predict ──────────────────────────────────────────────
def predict(text: str) -> dict:
    cleaned  = clean_text(text)
    encoding = tokenizer(
        cleaned,
        max_length=MAX_LEN,
        padding="max_length",
        truncation=True,
        return_tensors="pt"
    )
    input_ids      = encoding["input_ids"].to(DEVICE)
    attention_mask = encoding["attention_mask"].to(DEVICE)

    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)

    probabilities               = torch.softmax(outputs.logits, dim=1)
    confidence, predicted_class = torch.max(probabilities, dim=1)
    predicted_label             = label_encoder.inverse_transform([predicted_class.item()])[0]
    confidence_score            = round(confidence.item() * 100, 2)

    top3_probs, top3_indices = torch.topk(probabilities, 3, dim=1)
    top3 = [
        {
            "intent":     label_encoder.inverse_transform([idx.item()])[0],
            "confidence": round(prob.item() * 100, 2)
        }
        for prob, idx in zip(top3_probs[0], top3_indices[0])
    ]

    is_low_confidence = confidence_score < CONFIDENCE_THRESHOLD
    if predicted_label == OOS_INTENT or is_low_confidence:
        predicted_label = "out_of_scope"
        is_oos = True
    else:
        is_oos = False

    return {
        "input_text":        text,
        "cleaned_text":      cleaned,
        "intent":            predicted_label,
        "confidence":        confidence_score,
        "top3":              top3,
        "is_low_confidence": is_low_confidence,
        "is_oos":            is_oos,
        "response":          OOS_RESPONSE if is_oos else f"I detected your intent as {predicted_label}"
    }

# ── FastAPI App ──────────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    print("✅ API started!")
    yield

app = FastAPI(
    title="Intent Classification Chatbot API",
    version="2.0.0",
    lifespan=lifespan
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"]
)

# ── Models ───────────────────────────────────────────────
class PredictRequest(BaseModel):
    text:       str
    session_id: Optional[str] = None

class IntentResult(BaseModel):
    intent:     str
    confidence: float

class PredictResponse(BaseModel):
    session_id:        str
    input_text:        str
    cleaned_text:      str
    intent:            str
    confidence:        float
    top3:              list[IntentResult]
    last_intent:       Optional[str] = None
    is_low_confidence: bool = False
    is_oos:            bool = False
    response:          str = ""

# ── Routes ───────────────────────────────────────────────
@app.get("/")
def root():
    return {"message": "Intent Chatbot API", "status": "running", "intents": len(label_encoder.classes_)}

@app.get("/health")
def health():
    return {"status": "healthy", "model": "bert-base-uncased", "intents": len(label_encoder.classes_)}

@app.post("/predict", response_model=PredictResponse)
def predict_intent(request: PredictRequest):
    text = request.text.strip()
    if not text:
        raise HTTPException(status_code=400, detail="Text cannot be empty")
    if len(text) < 2:
        raise HTTPException(status_code=400, detail="Text too short")
    if len(text) > 200:
        raise HTTPException(status_code=400, detail="Text too long")

    import uuid, json
    from datetime import datetime

    session_id = request.session_id or str(uuid.uuid4())

    # Get last intent
    existing = r.get(f"session:{session_id}")
    messages = json.loads(existing) if existing else []
    last_intent = None
    for msg in reversed(messages):
        if msg.get("intent"):
            last_intent = msg["intent"]
            break

    result = predict(text)

    # Save to Redis
    messages.append({"role": "user", "text": text, "timestamp": datetime.now().isoformat()})
    messages.append({"role": "bot", "text": result["intent"], "intent": result["intent"], "confidence": result["confidence"], "timestamp": datetime.now().isoformat()})
    r.setex(f"session:{session_id}", 3600, json.dumps(messages))

    return PredictResponse(
        session_id=session_id,
        input_text=result["input_text"],
        cleaned_text=result["cleaned_text"],
        intent=result["intent"],
        confidence=result["confidence"],
        top3=[IntentResult(**r_) for r_ in result["top3"]],
        last_intent=last_intent,
        is_low_confidence=result["is_low_confidence"],
        is_oos=result["is_oos"],
        response=result["response"]
    )

@app.get("/history/{session_id}")
def get_history(session_id: str):
    existing = r.get(f"session:{session_id}")
    if not existing:
        raise HTTPException(status_code=404, detail="Session not found")
    return {"session_id": session_id, "history": json.loads(existing)}

@app.delete("/history/{session_id}")
def delete_session(session_id: str):
    r.delete(f"session:{session_id}")
    return {"message": "Session cleared"}

@app.get("/intents")
def get_intents():
    return {"total": len(label_encoder.classes_), "intents": list(label_encoder.classes_)}