import os
import io
import time
import json
import base64
import re
from typing import Optional, List
from fastapi import APIRouter, Depends, HTTPException, Form, UploadFile, File, Request
from fastapi.responses import JSONResponse
from google import genai
from google.genai import types
from PIL import Image

from schemas.user import UserBase
from utils.jwt import verify_token
from re_usable_function_tts import generate_tts_base64_audio
from gemini_tts import gemini_tts

router = APIRouter(prefix="/dictation", tags=["dictation"])

async def get_dictation_user(request: Request):
    """
    Extracts and verifies the token from cookies or Authorization header.
    """
    token = None
    auth_header = request.headers.get("Authorization")
    if auth_header and auth_header.startswith("Bearer "):
        token = auth_header.split(" ")[1]
    
    if not token:
        token = request.cookies.get("access_token")
    
    if not token:
        raise HTTPException(status_code=401, detail="Not authenticated")
    
    payload = verify_token(token)
    if not payload:
        raise HTTPException(status_code=401, detail="Invalid token")
    
    return UserBase(email=payload.get("email"), name=payload.get("name"))

GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
client = genai.Client(api_key=GOOGLE_API_KEY)
LLM_MODEL = 'gemini-2.5-flash' # Match main.py for compatibility

# In-memory session storage for development
dictation_sessions = {}

@router.post("/process-content")
async def process_content(
    request: Request,
    content_type: str = Form(...), # "text", "image", "pdf"
    text: Optional[str] = Form(None),
    file: Optional[UploadFile] = File(None),
    current_user: UserBase = Depends(get_dictation_user)
):
    """
    Extracts text from various inputs using Gemini 1.5 Flash.
    """
    try:
        if content_type == "text":
            return {"text": text}
        
        if not file:
            raise HTTPException(status_code=400, detail="File is required for image/pdf content type")

        file_bytes = await file.read()
        mime_type = file.content_type

        try:
            uploaded_file = client.files.upload(
                file=io.BytesIO(file_bytes),
                config={'mime_type': mime_type}
            )

            prompt = "Extract all text from this content. Return only the extracted text, cleaned and ready to be read aloud. No markdown, no extra commentary."
            response = client.models.generate_content(
                model=LLM_MODEL,
                contents=[uploaded_file, prompt]
            )
            
            return {"text": response.text}
        except Exception as gem_err:
            print(f"Gemini OCR error: {gem_err}")
            raise HTTPException(status_code=500, detail="AI text extraction failed.")

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Content processing failed: {str(e)}")

@router.post("/start")
async def start_dictation(
    request: Request,
    client_id: str = Form(...),
    text: str = Form(...),
    language: str = Form("english"),
    speed: str = Form("slow"), 
    chunk_mode: str = Form("sentence"),
    voice: str = Form("male"),
    current_user: UserBase = Depends(get_dictation_user)
):
    """
    Initializes a dictation session, chunks the text, and provides the first chunk metadata.
    """
    # 1. Advanced Chunking Logic
    if chunk_mode == "word":
        chunks = text.split()
    elif chunk_mode == "phrase":
        chunks = [p.strip() for p in re.split(r'[,;:]', text) if p.strip()]
    else: # sentence
        chunks = [s.strip() for s in re.split(r'[.!?]', text) if s.strip()]

    if not chunks:
        # Fallback if no punctuation found
        chunks = [text]

    # 2. Store session
    dictation_sessions[client_id] = {
        "text": text,
        "chunks": chunks,
        "current_index": 0,
        "settings": {
            "language": language,
            "speed": speed,
            "chunk_mode": chunk_mode,
            "voice": voice
        },
        "status": "in_progress",
        "created_at": time.time()
    }

    return {
        "message": "Dictation session started",
        "total_chunks": len(chunks),
        "current_index": 0,
        "welcome_message": f"Hello {current_user.name if hasattr(current_user, 'name') else 'Student'}. Today we will do dictation. Don’t worry. Write slowly. You can improve."
    }

@router.post("/get-chunk")
async def get_chunk(
    client_id: str = Form(...),
    index: int = Form(...),
    current_user: UserBase = Depends(get_dictation_user)
):
    """
    Returns the text and generated audio for a specific chunk.
    """
    session = dictation_sessions.get(client_id)
    if not session:
        raise HTTPException(status_code=404, detail="Session not found or expired.")

    if index < 0 or index >= len(session["chunks"]):
        raise HTTPException(status_code=400, detail="Invalid chunk index.")

    chunk_text = session["chunks"][index]
    language = session["settings"]["language"]
    voice = session["settings"]["voice"]
    
    session["current_index"] = index

    audio_base64 = None
    try:
        # Priority 1: Gemini TTS (Supports multiple languages well)
        if language.lower() in ["english", "malayalam", "hindi", "urdu"]:
            audio_data = gemini_tts(content=chunk_text, language=language)
            if audio_data:
                audio_base64 = base64.b64encode(audio_data).decode('utf-8')
        
        # Priority 2: Google Cloud TTS (Wavenet/Neural2)
        if not audio_base64:
            audio_base64 = generate_tts_base64_audio(text_to_speak=chunk_text, voice_name=voice, language=language)
            
    except Exception as e:
        print(f"TTS Error in dictation: {e}")

    return {
        "chunk_text": chunk_text,
        "audio_base64": audio_base64,
        "index": index,
        "is_last": index == len(session["chunks"]) - 1
    }

@router.post("/complete")
async def complete_dictation(
    client_id: str = Form(...),
    current_user: UserBase = Depends(get_dictation_user)
):
    """
    Finalizes the dictation activity but keeps the session for Comparison Mode.
    """
    session = dictation_sessions.get(client_id)
    if not session:
        return {"message": "Session not found."}
    
    session["status"] = "completed"
    return {
        "message": "Good job. Now let’s check together.",
        "original_text": session["text"]
    }

@router.post("/evaluate")
async def evaluate_dictation(
    client_id: str = Form(...),
    file: UploadFile = File(...),
    current_user: UserBase = Depends(get_dictation_user)
):
    """
    OCR student's handwritten work, compare it with original, and score.
    """
    session = dictation_sessions.get(client_id)
    if not session:
        # Fallback: maybe session expired, but we can still try to compare if we have the text
        # For now, let's require a session.
        raise HTTPException(status_code=404, detail="Session expired or not found. Please restart.")

    original_text = session["text"]
    file_bytes = await file.read()

    try:
        # 1. Upload handwritten work to Gemini
        uploaded_file = client.files.upload(
            file=io.BytesIO(file_bytes),
            config={'mime_type': file.content_type}
        )

        # 2. Combined OCR and Comparison Prompt
        prompt = f"""
        You are an expert supportive teacher.
        Student's Goal: Write the following original text accurately by hand from dictation.
        Original Text: "{original_text}"

        Your Task:
        1. Transcribe the student's handwritten work from the attached image.
        2. Compare it with the Original Text.
        3. Identify any spelling mistakes, missing words, or punctuation errors.
        4. Calculate a Score out of 100 (Be fair but accurate).
        5. Provide encouraging, supportive feedback in a friendly tone.

        Return ONLY a JSON object with this exact structure:
        {{
            "score": <number>,
            "transcribed_text": "<the text you saw in the image>",
            "mistakes": ["list", "of", "specific", "errors"],
            "feedback": "<your supportive feedback>",
            "original_text": "{original_text}"
        }}
        """

        response = client.models.generate_content(
            model=LLM_MODEL,
            contents=[uploaded_file, prompt],
            config=types.GenerateContentConfig(response_mime_type="application/json")
        )

        result = json.loads(response.text)
        
        # Cleanup session
        dictation_sessions.pop(client_id, None)

        return result

    except Exception as e:
        print(f"Evaluation Error: {e}")
        raise HTTPException(status_code=500, detail=f"Evaluation failed: {str(e)}")
