import os

from dotenv import load_dotenv
from typing import List, Optional, Dict
from pydantic import BaseModel
from fastapi import Body
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from utils.ai import AiService
from utils.vector_db import VectorDB
from lightrag import LightRAG, QueryParam
from lightrag.llm.openai import gpt_4_1_mini_complete, openai_embed
from utils.time_funcs import extract_timestamps
from utils.document_management import load_document, transcript_chunker, merge_chunks

from prompts import name_generation_prompt

# Load environment variables
load_dotenv()

# Initialize FastAPI app
app = FastAPI(title="Transcription API")

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

# Initialize services
working_dir = os.path.join(os.getcwd(), "_backend", "data","honduras_data")
if os.getenv("ENVIRONMENT") == "production":
    working_dir = os.path.join(os.getcwd(), "data","honduras_data")

transcription_data_dir = os.path.join(working_dir)
vector_db_name = "transcription"
vector_db_embedding_length = 1536

vector_db = VectorDB(
    vector_db_name,
    vector_db_embedding_length,
    os.getenv("OPENAI_API_KEY"),
    working_dir
)

ai_service = AiService(model="gpt-4.1")

lightrag_dir = os.path.join(working_dir, "lightrag")
rag = LightRAG(
    working_dir=lightrag_dir,
    llm_model_func=gpt_4_1_mini_complete,
    embedding_func=openai_embed
)

class ChatRequest(BaseModel):
    message: str
    chat_history: List[Dict[str, str]]


class ChatResponse(BaseModel):
    response: str
    timestamps: List[Dict[str, str | int]]  # Updated type annotation


class NameRequest(BaseModel):
    user_query: str
    assistant_response: str


class SearchRequest(BaseModel):
    query: str
    filters: Optional[Dict] = {}


class SessionInfoResponse(BaseModel):
    session_info: Dict
    video_url: str


# API Endpoints
@app.post("/api/v1/chat", response_model=ChatResponse)
async def chat_endpoint(request: ChatRequest):
    try:
        
        translated_query = ai_service.generate_single_response(
            request.message,
            "Translate the following text to Spanish",
            json_output=False
        )
        # Search documents
        search_docs = vector_db.search(translated_query, limit=100,
                                        # filter=filters,
                                        score_threshold=0
                                        )

        # Get GraphML response using asyncio.create_task to avoid event loop issues
        import asyncio
        loop = asyncio.get_event_loop()
        graphML_resp = await loop.run_in_executor(
            None,
            lambda: rag.query(
                request.message,
                param=QueryParam(
                    mode="hybrid",
                )
            )
        )

        # Generate response
        assistant_response = ai_service.generate_response(
            request.chat_history,
            docs=search_docs,
            extra_context=graphML_resp,
        )

        # Extract timestamps
        timestamps = extract_timestamps(assistant_response, search_docs)

        return ChatResponse(
            response=assistant_response,
            timestamps=timestamps  # Remove json.dumps(), return list directly
        )
    except Exception as e:
        print(e)
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/v1/generate_name", response_model=str)
async def generate_name_endpoint(request: NameRequest):
    try:
        messages = ai_service.init_messages(name_generation_prompt)
        ai_service.add_user_message(messages, request.user_query)

        # Generate response
        response = ai_service.generate_response(messages)

        return response
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/v1/transcript", response_model=List[Dict])
async def get_transcript_file(body: Dict = Body(...)):
    try:
        file_name = body.get("file_name")
        if not file_name:
            raise HTTPException(status_code=400, detail="file_name is required")
        file_path = os.path.join(transcription_data_dir, file_name)
        doc = load_document(file_path)
        transcript = transcript_chunker(doc, filename=file_name)
        transcript = merge_chunks(transcript)
        return transcript

    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="File not found")


@app.get("/api/v1/health")
async def health_check():
    return {"status": "healthy"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)