# =========================================
# 1. IMPORTS
# =========================================
import asyncio
import os
import json
import uuid
import cloudinary
import cloudinary.uploader
import firebase_admin

from firebase_admin import credentials, firestore
from fastapi import FastAPI, HTTPException, BackgroundTasks
from pydantic import BaseModel
from gradio_client import Client
from google.cloud.firestore_v1.base_query import FieldFilter
import edge_tts

from typing import Optional, List
from dotenv import load_dotenv
from contextlib import asynccontextmanager

# =========================================
# 2. INITIALIZATIONS & CONFIG
# =========================================
load_dotenv()

if not firebase_admin._apps:
    fb_json = os.getenv("FIREBASE_JSON")

    if fb_json:
        cred_dict = json.loads(fb_json)
        cred = credentials.Certificate(cred_dict)
    else:
        cred = credentials.Certificate("serviceAccountKey.json")

    firebase_admin.initialize_app(cred)

db = firestore.client()

cloudinary.config(
    cloud_name=os.getenv("CLOUD_NAME"),
    api_key=os.getenv("API_KEY"),
    api_secret=os.getenv("API_SECRET"),
    secure=True
)

HF_SPACE = "Fayza38/Question_and_answer_model"
client = None

# =========================================
# 3. CONSTANTS & MODELS
# =========================================
TECH_CATEGORIES = {
    0: "Security",
    1: "BackEnd",
    2: "Networking",
    3: "FrontEnd",
    4: "DataEngineering",
    5: "WebDevelopment",
    6: "FullStack",
    7: "VersionControl",
    8: "SystemDesign",
    9: "MachineLearning",
    10: "LanguagesAndFrameworks",
    11: "DatabaseSystems",
    12: "ArtificialIntelligence",
    13: "SoftwareTesting",
    14: "DistributedSystems",
    15: "DevOps",
    16: "LowLevelSystems",
    17: "DatabaseAndSql",
    18: "GeneralProgramming",
    19: "DataStructures",
    20: "Algorithms"
}

DIFFICULTY_MAP = {
    0: "Easy",
    1: "Intermediate",
    2: "Hard"
}


class GenerateSessionRequest(BaseModel):
    sessionId: str

    # 0 = Behavioral
    # 1 = Technical
    sessionType: int

    # IMPORTANT:
    # difficultyLevel is ONLY used for technical sessions
    difficultyLevel: Optional[int] = None

    # ONLY required for technical sessions
    trackName: Optional[int] = None


class CleanupRequest(BaseModel):
    audioUrls: List[str]


# =========================================
# 4. LIFESPAN MANAGEMENT
# =========================================
@asynccontextmanager
async def lifespan(app: FastAPI):
    global client

    print("Connecting to Hugging Face Model...")

    try:
        loop = asyncio.get_event_loop()

        client = await loop.run_in_executor(
            None,
            lambda: Client(HF_SPACE)
        )

        print("Model Connected Successfully!")

    except Exception as e:
        print(f"Model Connection Failed: {e}")

    yield

    print("Shutting down Intervision Service...")


app = FastAPI(
    title="Intervision AI Question Service",
    lifespan=lifespan
)

# =========================================
# 5. HELPER FUNCTIONS
# =========================================
async def generate_audio(text: str, filename: str):
    try:
        communicate = edge_tts.Communicate(
            text,
            "en-US-GuyNeural",
            rate="-15%"
        )

        await communicate.save(filename)

        upload_result = cloudinary.uploader.upload(
            filename,
            resource_type="video",
            folder="interview_audio"
        )

        if os.path.exists(filename):
            os.remove(filename)

        return upload_result["secure_url"]

    except Exception as e:
        print(f"Audio Generation Error: {e}")

        if os.path.exists(filename):
            os.remove(filename)

        return None


async def safe_generate(prompt: str, retries: int = 5):
    if client is None:
        raise Exception("AI Client is not initialized.")

    for attempt in range(retries):
        try:
            loop = asyncio.get_running_loop()

            result = await loop.run_in_executor(
                None,
                lambda: client.predict(
                    prompt=prompt,
                    api_name="/generate_questions"
                )
            )

            return result

        except Exception as e:
            if attempt == retries - 1:
                raise e

            await asyncio.sleep(5)


def parse_question_output(raw_output: str):
    if not raw_output:
        return None, None

    text = (
        raw_output.split("assistant")[-1].strip()
        if "assistant" in raw_output
        else raw_output
    )

    if "Q:" in text and "A:" in text:
        try:
            parts = text.split("A:")

            question = parts[0].replace("Q:", "").strip()

            answer = (
                parts[1]
                .split("<|im_end|>")[0]
                .strip()
            )

            return question, answer

        except Exception:
            return None, None

    return None, None


# =========================================
# 6. REFILL QUESTION POOLS
# =========================================
async def refill_specific_pool(
    track_id: Optional[int],
    difficulty: Optional[int],
    count: int,
    session_type: int
):
    while client is None:
        await asyncio.sleep(5)

    # =====================================
    # Behavioral Questions
    # =====================================
    if session_type == 0:

        track_text = "Behavioral"

        prompt = (
            "Generate ONE unique behavioral interview question. "
            "Format: Q: [Question] A: [Answer]"
        )

    # =====================================
    # Technical Questions
    # =====================================
    else:

        track_text = TECH_CATEGORIES.get(track_id)

        level_text = DIFFICULTY_MAP.get(difficulty)

        prompt = (
            f"Generate ONE unique {track_text} "
            f"question for {level_text} level. "
            f"Format: Q: [Question] A: [Answer]"
        )

    success_count = 0

    while success_count < count:
        try:
            raw_output = await safe_generate(prompt)

            question, answer = parse_question_output(raw_output)

            if question and answer:

                audio_url = await generate_audio(
                    question,
                    f"{uuid.uuid4()}.mp3"
                )

                if audio_url:

                    question_data = {
                        "session_type": session_type,
                        "questionText": question,
                        "questionIdealAnswer": answer,
                        "audio_url": audio_url,
                        "created_at": firestore.SERVER_TIMESTAMP
                    }

                    # =================================
                    # Technical ONLY
                    # =================================
                    if session_type == 1:
                        question_data["track_id"] = track_id
                        question_data["difficulty"] = difficulty

                    db.collection("questions_pool").add(question_data)

                    success_count += 1

                    print(
                        f"Successfully added "
                        f"{track_text} question "
                        f"{success_count}/{count}"
                    )

                    await asyncio.sleep(2)

        except Exception as e:
            print(f"Refill error: {e}")
            await asyncio.sleep(5)

# =========================================
# 7. API ENDPOINTS
# =========================================
@app.post("/generate-session")
async def generate_session(
    request: GenerateSessionRequest,
    background_tasks: BackgroundTasks
):
    session_type = request.sessionType
    track_id = request.trackName

    # =====================================
    # Behavioral Session
    # =====================================
    if session_type == 0:

        query = db.collection("questions_pool").where(
            filter=FieldFilter("session_type", "==", 0)
        )

    # =====================================
    # Technical Session
    # =====================================
    elif session_type == 1:

        if track_id is None:
            raise HTTPException(
                status_code=400,
                detail="trackName is required for technical sessions."
            )

        if request.difficultyLevel is None:
            raise HTTPException(
                status_code=400,
                detail="difficultyLevel is required for technical sessions."
            )

        difficulty = request.difficultyLevel

        query = (
            db.collection("questions_pool")
            .where(filter=FieldFilter("session_type", "==", 1))
            .where(filter=FieldFilter("track_id", "==", track_id))
            .where(filter=FieldFilter("difficulty", "==", difficulty))
        )

    else:
        raise HTTPException(
            status_code=400,
            detail="Invalid sessionType."
        )

    docs = query.limit(10).get()

    final_questions = []

    for index, doc in enumerate(docs, start=1):

        data = doc.to_dict()

        final_questions.append({
            "question_id": index,
            "text": data["questionText"],
            "expected_answer": data["questionIdealAnswer"],
            "audio_url": data.get("audio_url", "")
        })

        # remove used question
        db.collection("questions_pool").document(doc.id).delete()

    # =====================================
    # BACKGROUND REFILL
    # =====================================
    async def check_and_refill_background():

        snap = query.count().get()
        current_count = snap[0][0].value

        if current_count < 50:

            if session_type == 0:

                print(
                    f"Behavioral stock low "
                    f"({current_count}) -> refilling..."
                )

                await refill_specific_pool(
                    track_id=None,
                    difficulty=None,
                    count=50 - current_count,
                    session_type=0
                )

            else:

                print(
                    f"{TECH_CATEGORIES[track_id]} stock low "
                    f"({current_count}) -> refilling..."
                )

                await refill_specific_pool(
                    track_id=track_id,
                    difficulty=difficulty,
                    count=50 - current_count,
                    session_type=1
                )

    background_tasks.add_task(check_and_refill_background)

    if not final_questions:
        raise HTTPException(
            status_code=503,
            detail="The question pool is empty. Please try again in a few minutes."
        )

    return {
        "session_id": request.sessionId,
        "questions": final_questions
    }


# =========================================
# 8. ADMIN PREFILL
# =========================================
@app.get("/admin/prefill-all")
async def prefill_all(background_tasks: BackgroundTasks):

    async def run_sync():

        print("Starting Global Smart Prefill...")

        # =================================
        # Behavioral Questions
        # =================================
        behavioral_query = (
            db.collection("questions_pool")
            .where(filter=FieldFilter("session_type", "==", 0))
        )

        behavioral_snap = behavioral_query.count().get()

        behavioral_count = behavioral_snap[0][0].value

        if behavioral_count < 50:

            needed = 50 - behavioral_count

            print(f"Syncing Behavioral: adding {needed}")

            await refill_specific_pool(
                track_id=None,
                difficulty=None,
                count=needed,
                session_type=0
            )

        # =================================
        # Technical Questions
        # =================================
        for track_id, track_name in TECH_CATEGORIES.items():

            for diff_id, diff_name in DIFFICULTY_MAP.items():

                technical_query = (
                    db.collection("questions_pool")
                    .where(filter=FieldFilter("session_type", "==", 1))
                    .where(filter=FieldFilter("track_id", "==", track_id))
                    .where(filter=FieldFilter("difficulty", "==", diff_id))
                )

                snap = technical_query.count().get()

                current = snap[0][0].value

                if current < 50:

                    needed = 50 - current

                    print(
                        f"Syncing {track_name} "
                        f"({diff_name}): adding {needed}"
                    )

                    await refill_specific_pool(
                        track_id=track_id,
                        difficulty=diff_id,
                        count=needed,
                        session_type=1
                    )

        print("Global Smart Prefill Completed!")

    background_tasks.add_task(run_sync)

    return {
        "message": "Global prefill process started in the background."
    }


# =========================================
# 9. CLEANUP AUDIO
# =========================================
@app.post("/cleanup-audio")
async def cleanup_audio(
    request: CleanupRequest,
    background_tasks: BackgroundTasks
):
    def delete_job(urls):

        for url in urls:
            try:
                public_id = (
                    "interview_audio/"
                    + url.split("/")[-1].split(".")[0]
                )

                cloudinary.uploader.destroy(
                    public_id,
                    resource_type="video"
                )

            except Exception:
                pass

    background_tasks.add_task(
        delete_job,
        request.audioUrls
    )

    return {
        "message": "Cloudinary cleanup process initiated."
    }


# =========================================
# 10. HEALTH CHECK
# =========================================
@app.get("/health")
async def health():
    return {
        "status": "active",
        "ai_model_connected": client is not None
    }


@app.get("/")
async def root():
    return {
        "app": "Intervision AI Engine",
        "status": "Running"
    }


# =========================================
# 11. RUN SERVER
# =========================================
if __name__ == "__main__":

    import uvicorn

    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=True
    )