from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import torch
from transformers import pipeline
import json
import os
from difflib import SequenceMatcher
from typing import Dict, Any, Optional
import tempfile
import subprocess
import shutil

app = FastAPI(
    title="Bayan AI بيان",
    description="",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allow all origins for local development
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# CPU only on free tier
device = -1

# Load Whisper pipeline (model downloads/caches automatically on first run)
pipe = pipeline(
    "automatic-speech-recognition",
    model="tarteel-ai/whisper-tiny-ar-quran",
    device=device,
)

# Standard Surah names (1–114)
surah_names = {
    1: "Al-Fatiha (الفاتحة)",
    2: "Al-Baqarah (البقرة)",
    3: "Aal-E-Imran (آل عمران)",
    4: "An-Nisa (النساء)",
    5: "Al-Maidah (المائدة)",
    6: "Al-An'am (الأنعام)",
    7: "Al-A'raf (الأعراف)",
    8: "Al-Anfal (الأنفال)",
    9: "At-Tawbah (التوبة)",
    10: "Yunus (يونس)",
    11: "Hud (هود)",
    12: "Yusuf (يوسف)",
    13: "Ar-Ra'd (الرعد)",
    14: "Ibrahim (إبراهيم)",
    15: "Al-Hijr (الحجر)",
    16: "An-Nahl (النحل)",
    17: "Al-Isra (الإسراء)",
    18: "Al-Kahf (الكهف)",
    19: "Maryam (مريم)",
    20: "Ta-Ha (طه)",
    21: "Al-Anbiya (الأنبياء)",
    22: "Al-Hajj (الحج)",
    23: "Al-Mu'minun (المؤمنون)",
    24: "An-Nur (النور)",
    25: "Al-Furqan (الفرقان)",
    26: "Ash-Shu'ara (الشعراء)",
    27: "An-Naml (النمل)",
    28: "Al-Qasas (القصص)",
    29: "Al-Ankabut (العنكبوت)",
    30: "Ar-Rum (الروم)",
    31: "Luqman (لقمان)",
    32: "As-Sajdah (السجدة)",
    33: "Al-Ahzab (الأحزاب)",
    34: "Saba (سبأ)",
    35: "Fatir (فاطر)",
    36: "Ya-Sin (يس)",
    37: "As-Saffat (الصافات)",
    38: "Sad (ص)",
    39: "Az-Zumar (الزمر)",
    40: "Ghafir (غافر)",
    41: "Fussilat (فصلت)",
    42: "Ash-Shura (الشورى)",
    43: "Az-Zukhruf (الزخرف)",
    44: "Ad-Dukhkhan (الدخان)",
    45: "Al-Jathiya (الجاثية)",
    46: "Al-Ahqaf (الأحقاف)",
    47: "Muhammad (محمد)",
    48: "Al-Fath (الفتح)",
    49: "Al-Hujurat (الحجرات)",
    50: "Qaf (ق)",
    51: "Adh-Dhariyat (الذاريات)",
    52: "At-Tur (الطور)",
    53: "An-Najm (النجم)",
    54: "Al-Qamar (القمر)",
    55: "Ar-Rahman (الرحمن)",
    56: "Al-Waqi'ah (الواقعة)",
    57: "Al-Hadid (الحديد)",
    58: "Al-Mujadila (المجادلة)",
    59: "Al-Hashr (الحشر)",
    60: "Al-Mumtahina (الممتحنة)",
    61: "As-Saff (الصف)",
    62: "Al-Jumu'ah (الجمعة)",
    63: "Al-Munafiqoon (المنافقون)",
    64: "At-Taghabun (التغابن)",
    65: "At-Talaq (الطلاق)",
    66: "At-Tahrim (التحريم)",
    67: "Al-Mulk (الملك)",
    68: "Al-Qalam (القلم)",
    69: "Al-Haqqah (الحاقة)",
    70: "Al-Ma'arij (المعارج)",
    71: "Nooh (نوح)",
    72: "Al-Jinn (الجن)",
    73: "Al-Muzzammil (المزمل)",
    74: "Al-Muddathir (المدثر)",
    75: "Al-Qiyamah (القيامة)",
    76: "Al-Insan (الإنسان)",
    77: "Al-Mursalat (المرسلات)",
    78: "An-Naba (النبأ)",
    79: "An-Nazi'at (النازعات)",
    80: "Abasa (عبس)",
    81: "At-Takwir (التكوير)",
    82: "Al-Infitar (الإنفطار)",
    83: "Al-Mutaffifin (المطففين)",
    84: "Al-Inshiqaq (الإنشقاق)",
    85: "Al-Buruj (البروج)",
    86: "At-Tariq (الطارق)",
    87: "Al-A'la (الأعلى)",
    88: "Al-Ghashiyah (الغاشية)",
    89: "Al-Fajr (الفجر)",
    90: "Al-Balad (البلد)",
    91: "Ash-Shams (الشمس)",
    92: "Al-Lail (الليل)",
    93: "Ad-Duha (الضحى)",
    94: "Ash-Sharh (الشرح)",
    95: "At-Tin (التين)",
    96: "Al-Alaq (العلق)",
    97: "Al-Qadr (القدر)",
    98: "Al-Bayyina (البينة)",
    99: "Az-Zalzalah (الزلزلة)",
    100: "Al-Adiyat (العاديات)",
    101: "Al-Qari'ah (القارعة)",
    102: "At-Takathur (التكاثر)",
    103: "Al-Asr (العصر)",
    104: "Al-Humazah (الهمزة)",
    105: "Al-Fil (الفيل)",
    106: "Quraish (قريش)",
    107: "Al-Ma'un (الماعون)",
    108: "Al-Kawthar (الكوثر)",
    109: "Al-Kafirun (الكافرون)",
    110: "An-Nasr (النصر)",
    111: "Al-Masad (المسد)",
    112: "Al-Ikhlas (الإخلاص)",
    113: "Al-Falaq (الفلق)",
    114: "An-Nas (الناس)",
}

# Phrases to ignore (e.g., common introductions)
PHRASES_TO_IGNORE = [
    "بِسْمِ اللَّهِ الرَّحْمَنِ الرَّحِيمِ",
    "أعوذ بالله من الشيطان الرجيم",
    "صدق الله العظيم",
]

import re

def normalize_text(text: str) -> str:
    """Robust normalization for Arabic text."""
    text = re.sub(r"[إأآاٱ]", "ا", text)
    text = re.sub(r"ى", "ي", text)
    text = re.sub(r"ؤ", "ء", text)
    text = re.sub(r"ئ", "ء", text)
    text = re.sub(r"g", "ة", text)
    text = re.sub(r"ة", "ه", text)
    text = re.sub(r"[\u064B-\u065F\u0670]", "", text) # Tashkeel
    text = re.sub(r"[\u06D6-\u06ED]", "", text) 
    text = re.sub(r"ء", "", text) # Remove Hamza to handle varying forms
    return " ".join(text.strip().split())

# Pre-load all verses at startup
all_verses = []

surahs_dir = "surahs_json_files"
if not os.path.isdir(surahs_dir):
    raise FileNotFoundError("Missing 'surahs_json_files/' folder.")

for filename in sorted(os.listdir(surahs_dir)):
    if filename.endswith(".json"):
        try:
            surah_number = int(filename.split("_")[0])
        except:
            continue
        surah_name = surah_names.get(surah_number, f"Surah {surah_number}")
        file_path = os.path.join(surahs_dir, filename)
        
        with open(file_path, "r", encoding="utf-8") as f:
            data = json.load(f)
        
        verses = [ayah["text"] for ayah in data.get("ayahs", []) if "text" in ayah]
        
        for ayah_number, verse_text in enumerate(verses, start=1):
            verse_norm = normalize_text(verse_text)
            all_verses.append({
                "surah_number": surah_number,
                "surah_name": surah_name,
                "ayah_number": ayah_number,
                "verse_text": verse_text,
                "verse_norm": verse_norm
            })

print(f"Loaded {len(all_verses)} verses from {len(os.listdir(surahs_dir))} surahs.")

def find_best_verse(transcription: str) -> Dict[str, Any]:
    transcription_norm = normalize_text(transcription)
    
    # Remove phrases to ignore (Bismillah, A'udhu billah)
    for phrase in PHRASES_TO_IGNORE:
        phrase_norm = normalize_text(phrase)
        if phrase_norm in transcription_norm:
            # Replace and clean up extra spaces
            transcription_norm = transcription_norm.replace(phrase_norm, "").strip()
            transcription_norm = " ".join(transcription_norm.split())

    if not transcription_norm:
        return {"error": "Empty transcription"}
    
    candidates = []
    
    # Pre-compile regex for whole word check
    pattern_str = r'(?:^|\s)' + re.escape(transcription_norm) + r'(?:\s|$)'
    whole_word_regex = re.compile(pattern_str)

    for verse in all_verses:
        verse_norm = verse["verse_norm"]
        
        is_whole_word = False
        containment = 0.0
        ratio = 0.0
        
        # Fast substring check
        if transcription_norm in verse_norm:
            containment = 1.0
            matcher = SequenceMatcher(None, transcription_norm, verse_norm)
            ratio = matcher.ratio()
            
            # Check for whole word match
            if whole_word_regex.search(verse_norm):
                is_whole_word = True
        else:
            matcher = SequenceMatcher(None, transcription_norm, verse_norm)
            match = matcher.find_longest_match(0, len(transcription_norm), 0, len(verse_norm))
            containment = match.size / len(transcription_norm) if len(transcription_norm) > 0 else 0
            ratio = matcher.ratio()
        
        candidates.append({
            "verse": verse,
            "containment": containment,
            "ratio": ratio,
            "is_whole_word": is_whole_word
        })
    
    # Sort by whole_word (desc), containment (desc), ratio (desc)
    candidates.sort(key=lambda x: (x["is_whole_word"], x["containment"], x["ratio"]), reverse=True)
    
    # If we have whole word matches, ignore partial matches
    if candidates and candidates[0]["is_whole_word"]:
        candidates = [c for c in candidates if c["is_whole_word"]]
    
    # Filter strong matches (>= 80% containment)
    strong_matches = [c for c in candidates if c["containment"] >= 0.8]
    
    def format_match(candidate):
        verse_data = candidate["verse"]
        return {
            "surah_number": verse_data["surah_number"],
            "surah_name": verse_data["surah_name"],
            "ayah_number": verse_data["ayah_number"],
            "verse_text": verse_data["verse_text"],
            "similarity_score": round(candidate["containment"], 4)
        }

    if not strong_matches:
        # No strong matches found
        if candidates:
            top_match = candidates[0]
            return {
                "error": "No confident match found",
                "best_similarity": round(top_match["containment"], 4),
                "possible_match": format_match(top_match)
            }
        else:
            return {"error": "No matches found"}

    if len(strong_matches) > 1:
        # Multiple strong matches -> return top 5
        top_5 = strong_matches[:5]
        return {
            "matches": [format_match(m) for m in top_5]
        }
    else:
        # Single dominant match
        return format_match(strong_matches[0])

@app.get("/")
def root():
    return {"message": "Bayan AI بيان... LIVE!"}

@app.post("/recognize")
async def recognize(file: UploadFile = File(...)):
    # Allow both audio and video
    is_video = file.content_type and file.content_type.startswith("video/")
    is_audio = file.content_type and file.content_type.startswith("audio/")
    
    if not is_audio and not is_video:
        raise HTTPException(status_code=400, detail="File must be an audio or video file")
    
    # Save to temp file
    contents = await file.read()
    file_extension = os.path.splitext(file.filename)[1] or (".mp4" if is_video else ".wav")
    
    with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp:
        tmp.write(contents)
        input_path = tmp.name
    
    audio_path = input_path
    temp_audio_path = None
    
    try:
        if is_video:
            # Check if ffmpeg is installed
            if not shutil.which("ffmpeg"):
                raise HTTPException(status_code=500, detail="ffmpeg not found on server")
            
            temp_audio_path = input_path + "_converted.wav"
            # Extract audio quickly and silently
            # -vn: no video, -acodec pcm_s16le: wav format, -ar 16000: whisper preferred sample rate
            # -y: overwrite, -loglevel error: be silent
            cmd = [
                "ffmpeg", "-y", "-i", input_path,
                "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1",
                "-loglevel", "error",
                temp_audio_path
            ]
            subprocess.run(cmd, check=True)
            audio_path = temp_audio_path
            
        transcription = pipe(audio_path)["text"]
    except subprocess.CalledProcessError as e:
        raise HTTPException(status_code=500, detail=f"Video conversion error: {str(e)}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Transcription error: {str(e)}")
    finally:
        # Clean up all temp files
        if os.path.exists(input_path):
            os.unlink(input_path)
        if temp_audio_path and os.path.exists(temp_audio_path):
            os.unlink(temp_audio_path)
    
    result = find_best_verse(transcription)
    result["transcription"] = transcription
    
    return JSONResponse(content=result)