import os, re, html, pickle
import numpy as np
import torch
from collections import Counter
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from openai import OpenAI

# ── Config ────────────────────────────────────────────────────────────────────
METER_MODEL_ID = "Rahaf2001/Lassen-meter-classifier"
ERA_MODEL_ID   = "Rahaf2001/LassenEraClassifier"
TOPIC_MODEL_ID = "Rahaf2001/Lassen-topic-classifier"

OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ── Global model holders ──────────────────────────────────────────────────────
models = {}

# ── Arabic text cleaning ──────────────────────────────────────────────────────
ARABIC_DIACRITICS = re.compile(r'[\u0617-\u061A\u064B-\u0652\u0670\u06D6-\u06ED]')

def clean_arabic(text: str) -> str:
    if not text:
        return ""
    text = html.unescape(str(text))
    text = re.sub(r"<.*?>", " ", text)
    text = text.replace("\u0640", "")
    text = ARABIC_DIACRITICS.sub("", text)
    text = re.sub(r'[\u0623\u0625\u0622\u0671]', '\u0627', text)
    text = text.replace("\u0629", "\u0647")
    text = re.sub(r"[0-9\u0660-\u0669]", " ", text)
    text = re.sub(r"[^\u0600-\u06FF\s]", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text

# ── Meter labels ──────────────────────────────────────────────────────────────
LABELS_METER = ['saree', 'kamel', 'mutakareb', 'mutadarak', 'munsareh',
                'madeed', 'mujtath', 'ramal', 'baseet', 'khafeef',
                'taweel', 'wafer', 'hazaj', 'rajaz']

METER_ARABIC = {
    'saree': 'السريع', 'kamel': 'الكامل', 'mutakareb': 'المتقارب',
    'mutadarak': 'المتدارك', 'munsareh': 'المنسرح', 'madeed': 'المديد',
    'mujtath': 'المجتث', 'ramal': 'الرمل', 'baseet': 'البسيط',
    'khafeef': 'الخفيف', 'taweel': 'الطويل', 'wafer': 'الوافر',
    'hazaj': 'الهزج', 'rajaz': 'الرجز'
}

# ── Meter taf'ila patterns ────────────────────────────────────────────────────
METER_PATTERNS = {
    'الطويل': 'فَعُولُنْ مَفَاعِيلُنْ فَعُولُنْ مَفَاعِلُنْ',
    'الكامل': 'مُتَفَاعِلُنْ مُتَفَاعِلُنْ مُتَفَاعِلُنْ',
    'البسيط': 'مُسْتَفْعِلُنْ فَاعِلُنْ مُسْتَفْعِلُنْ فَاعِلُنْ',
    'الوافر': 'مُفَاعَلَتُنْ مُفَاعَلَتُنْ فَعُولُنْ',
    'الخفيف': 'فَاعِلَاتُنْ مُسْتَفْعِلُنْ فَاعِلَاتُنْ',
    'الرجز': 'مُسْتَفْعِلُنْ مُسْتَفْعِلُنْ مُسْتَفْعِلُنْ',
    'الرمل': 'فَاعِلَاتُنْ فَاعِلَاتُنْ فَاعِلَاتُنْ',
    'السريع': 'مُسْتَفْعِلُنْ مُسْتَفْعِلُنْ مَفْعُولَاتُ',
    'المنسرح': 'مُسْتَفْعِلُنْ مَفْعُولَاتُ مُسْتَفْعِلُنْ',
    'الهزج': 'مَفَاعِيلُنْ مَفَاعِيلُنْ',
    'المتقارب': 'فَعُولُنْ فَعُولُنْ فَعُولُنْ فَعُولُنْ',
    'المتدارك': 'فَاعِلُنْ فَاعِلُنْ فَاعِلُنْ فَاعِلُنْ',
    'المديد': 'فَاعِلَاتُنْ فَاعِلُنْ فَاعِلَاتُنْ',
}

# ── Lifespan: load models once at startup ────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    print("Loading models...")

    # Meter
    models["meter_tokenizer"] = AutoTokenizer.from_pretrained(METER_MODEL_ID)
    models["meter_model"] = AutoModelForSequenceClassification.from_pretrained(METER_MODEL_ID)
    models["meter_model"].to(device).eval()

    # Era
    models["era_tokenizer"] = AutoTokenizer.from_pretrained(ERA_MODEL_ID)
    models["era_model"] = AutoModelForSequenceClassification.from_pretrained(ERA_MODEL_ID)
    models["era_model"].to(device).eval()

    # Topic
    models["topic_tokenizer"] = AutoTokenizer.from_pretrained(TOPIC_MODEL_ID)
    models["topic_model"] = AutoModelForSequenceClassification.from_pretrained(TOPIC_MODEL_ID)
    models["topic_model"].to(device).eval()

    # Topic labels — loaded from HF model config
    topic_cfg = models["topic_model"].config
    if hasattr(topic_cfg, "id2label"):
        models["id2label_topic"] = {int(k): v for k, v in topic_cfg.id2label.items()}
    else:
        models["id2label_topic"] = {i: str(i) for i in range(topic_cfg.num_labels)}

    # OpenAI client
    models["openai"] = OpenAI(api_key=OPENAI_API_KEY)

    print(f"All models loaded on {device} ✓")
    yield
    models.clear()

# ── App ───────────────────────────────────────────────────────────────────────
app = FastAPI(title="Bayan API", lifespan=lifespan)

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

# ── Inference helpers ─────────────────────────────────────────────────────────
def predict_meter(poem_text: str) -> dict:
    verses = [v.replace("#", " ").strip() for v in poem_text.strip().split("\n") if v.strip()]
    if not verses:
        raise ValueError("القصيدة فارغة")
    predictions = []
    for verse in verses:
        inputs = models["meter_tokenizer"](verse, return_tensors="pt", truncation=True,
                                           max_length=32, padding="max_length")
        inputs = {k: v.to(device) for k, v in inputs.items()}
        with torch.no_grad():
            probs = torch.softmax(models["meter_model"](**inputs).logits, dim=-1)[0]
        pred_id = torch.argmax(probs).item()
        predictions.append((LABELS_METER[pred_id], probs[pred_id].item()))
    top_meter = Counter(p[0] for p in predictions).most_common(1)[0][0]
    avg_conf = sum(c for _, c in predictions) / len(predictions)
    return {"meter_ar": METER_ARABIC[top_meter], "meter_en": top_meter,
            "confidence": round(avg_conf, 3)}

def predict_era(poem_text: str) -> dict:
    cleaned = clean_arabic(poem_text)
    enc = models["era_tokenizer"](cleaned, padding="max_length", truncation=True,
                                   max_length=256, return_tensors="pt")
    enc = {k: v.to(device) for k, v in enc.items()}
    with torch.no_grad():
        probs = torch.softmax(models["era_model"](**enc).logits, dim=-1).cpu().numpy()[0]
    label_names = ["قديم", "حديث"]
    pred_idx = int(np.argmax(probs))
    return {"era": label_names[pred_idx],
            "classical_probability": round(float(probs[0]), 4),
            "modern_probability": round(float(probs[1]), 4)}

def predict_topic(poem_text: str) -> dict:
    cleaned = clean_arabic(poem_text)
    inputs = models["topic_tokenizer"](cleaned, truncation=True, max_length=512,
                                        return_tensors="pt", padding=True)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.no_grad():
        probs = torch.softmax(models["topic_model"](**inputs).logits, dim=-1)[0].cpu().numpy()
    top3 = np.argsort(probs)[::-1][:3]
    id2label = models["id2label_topic"]
    return {"topic": id2label[int(top3[0])],
            "confidence": round(float(probs[top3[0]]), 3),
            "top3": [{"label": id2label[int(i)], "prob": round(float(probs[i]), 3)} for i in top3]}

# ── Request / Response schemas ────────────────────────────────────────────────
class PoemRequest(BaseModel):
    poem: str

class GenerateRequest(BaseModel):
    idea: str
    meter: str
    num_verses: int = 4

# ── Routes ────────────────────────────────────────────────────────────────────
@app.get("/")
def root():
    return {"status": "ok", "service": "Bayan API"}

@app.get("/health")
def health():
    return {"status": "healthy", "device": str(device)}

@app.post("/fasserha")
def fasserha(req: PoemRequest):
    """فسّرها لي — classify meter, era, topic then generate literary analysis."""
    if not req.poem.strip():
        raise HTTPException(400, "القصيدة فارغة")
    try:
        meter = predict_meter(req.poem)
        era   = predict_era(req.poem)
        topic = predict_topic(req.poem)
    except Exception as e:
        raise HTTPException(500, f"خطأ في التصنيف: {str(e)}")

    system_prompt = """أنت ناقد أدبي متخصص في الشعر العربي الكلاسيكي والحديث.
تحلل القصائد بأسلوب أكاديمي راقٍ، وتستخدم المصطلحات البلاغية والعروضية بدقة.
ردك دائماً بالعربية الفصحى."""

    user_prompt = f"""حلّل هذه القصيدة:

{req.poem}

معطيات النماذج (حقائق مؤكدة):
- البحر الشعري: {meter['meter_ar']} (ثقة: {meter['confidence']*100:.0f}%)
- العصر: {era['era']} (كلاسيكي: {era['classical_probability']*100:.0f}% | حديث: {era['modern_probability']*100:.0f}%)
- الموضوع: {topic['topic']} (ثقة: {topic['confidence']*100:.0f}%)

اكتب تحليلاً أدبياً شاملاً يتضمن:
1. الفكرة العامة والمعنى الكلي
2. المعنى التفصيلي للأبيات
3. الجماليات البلاغية والأسلوبية
4. البحر والإيقاع وأثرهما في المعنى
5. لمسة نقدية تقييمية

التزم بالترتيب أعلاه. لا تكرر المعلومات."""

    try:
        response = models["openai"].chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            max_tokens=1200,
            temperature=0.7
        )
        explanation = response.choices[0].message.content
    except Exception as e:
        raise HTTPException(500, f"خطأ في التفسير: {str(e)}")

    return {
        "success": True,
        "data": {
            "meter": meter,
            "era": era,
            "topic": topic,
            "explanation": explanation
        }
    }


@app.post("/generate")
def generate(req: GenerateRequest):
    """ساعدني في الكتابة — generate classical Arabic verses."""
    if not req.idea.strip():
        raise HTTPException(400, "الفكرة فارغة")
    if req.meter not in METER_PATTERNS and req.meter not in METER_ARABIC.values():
        raise HTTPException(400, f"البحر غير معروف: {req.meter}")
    if not 1 <= req.num_verses <= 12:
        raise HTTPException(400, "عدد الأبيات بين 1 و 12")

    pattern = METER_PATTERNS.get(req.meter, "")
    pattern_line = f"تفعيلة البحر: {pattern}" if pattern else ""

    prompt = f"""أنت شاعر عربي متخصص في العروض الكلاسيكي.

الموضوع: "{req.idea}"
البحر: {req.meter}
{pattern_line}

المطلوب: اكتب {req.num_verses} أبيات شعرية بالفصحى الكلاسيكية.

القواعد الصارمة:
- كل بيت من شطرين صحيحين عروضياً
- قافية موحدة في جميع الأبيات
- فصحى كلاسيكية فقط
- الأبيات متصلة كقصيدة واحدة
- اكتب الأبيات فقط، بدون ترقيم أو شرح
- سطر واحد لكل بيت، {req.num_verses} سطور فقط"""

    try:
        response = models["openai"].chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "أنت شاعر عربي كلاسيكي. اكتب الأبيات فقط، سطر لكل بيت."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.75,
            max_tokens=600
        )
        raw = response.choices[0].message.content.strip()
        verses = [l.strip() for l in raw.split("\n") if l.strip() and len(l.strip()) > 10]
        return {
            "success": True,
            "data": {
                "verses": verses[:req.num_verses],
                "meter": req.meter,
                "pattern": pattern
            }
        }
    except Exception as e:
        raise HTTPException(500, f"خطأ في التوليد: {str(e)}")


@app.get("/meters")
def list_meters():
    """Return all supported meters."""
    return {"meters": list(METER_PATTERNS.keys())}