import os
import shutil
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request, BackgroundTasks, logger
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, Response
from fastapi.staticfiles import StaticFiles
from contextlib import asynccontextmanager
from typing import Optional
import uvicorn
import base64
from pathlib import Path
import tempfile
import json
import pandas as pd
import logging
from curriculum_structure import convert_json_to_db_format
from process_pdf_pipline import run_full_pipeline

# Import your existing modules
from core import AppConfig
from repositories import MinIOStorageRepository
from services import (
    AudioService, ChatService, HealthService, ResponseService, 
    ResponseManager, OpenAIService, AgentService, ConnectionPool, LanguageSegmentationService,
    DataIngestionService
)



class DIContainer:
    def __init__(self):
        self.config = AppConfig.from_env()
        self.storage_repo = MinIOStorageRepository(self.config)
        self.response_manager = ResponseManager()
        
        # Initialize OpenAI and Agent services
        self.openai_service = OpenAIService()

        self.pool_handler = ConnectionPool(
            dbname=os.getenv("POSTGRES_DB"),
            user=os.getenv("POSTGRES_USER"),
            password=os.getenv("POSTGRES_PASSWORD"),
            host=os.getenv("DB_HOST"),
            port=int(os.getenv("DB_PORT"))
        )
        print(os.getenv("DB_HOST"), os.getenv("POSTGRES_DB"), os.getenv("POSTGRES_USER"))
        self.agent_service = AgentService(pool_handler=self.pool_handler)

        self.data_ingestion_service = DataIngestionService(pool_handler=self.pool_handler)


        # Initialize services
        self.audio_service = AudioService(self.storage_repo, self.config.minio_bucket)
        self.segmentation_service = LanguageSegmentationService()
        self.chat_service = ChatService(
            self.storage_repo, 
            self.response_manager, 
            self.config,
            self.openai_service,
            self.agent_service,
            self.segmentation_service
        )
        self.response_service = ResponseService(self.response_manager, self.audio_service)
        self.health_service = HealthService(self.storage_repo, self.config)


def run_processing_pipeline(pdf_path: str, grade: int, subject: str) -> tuple[str, str]:
    """
    Runs the full PDF processing pipeline and returns paths to the generated CSV and JSON files.
    """
    temp_json_path = "temp_json.json"
    temp_csv_path = "temp_embeddings.csv"
    run_full_pipeline(pdf_path, grade, subject, temp_json_path, temp_csv_path, remove_lessons=True)
    return temp_csv_path, temp_json_path


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Manages application startup and shutdown events for resource safety.
    """
    # --- Code to run ON STARTUP ---
    print("Application starting up...")
    container = DIContainer()
    app.state.container = container 
    print("DIContainer created and database pool initialized.")

    yield  # The application is now running and handling requests

    # --- Code to run ON SHUTDOWN ---
    print("Application shutting down...")
    # This is the guaranteed, graceful shutdown call
    app.state.container.agent_service.close() 
    print("Database connection pool closed successfully.")


def create_app() -> FastAPI:
    # Connect the lifespan manager to your FastAPI app instance
    app = FastAPI(title="Unified Chat API with Local Agent", lifespan=lifespan)
    logger = logging.getLogger("uvicorn.error")
    
    # Fixed CORS configuration for CapRover
    app.add_middleware(
        CORSMiddleware,
        allow_origins=[
            "https://voice-agent.caprover.al-arcade.com",
            "http://voice-agent.caprover.al-arcade.com",
            "http://localhost:8000",  # For local development
            "http://127.0.0.1:8000",
            "*"  # Allow all origins for testing - remove in production
        ],
        allow_credentials=True,
        allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
        allow_headers=[
            "Accept",
            "Accept-Language", 
            "Content-Language",
            "Content-Type",
            "Authorization",
            "X-Response-Text"
        ],
        expose_headers=["X-Response-Text"],
    )


    def process_pdf_curriculum_in_background(pdf_bytes: bytes, original_filename: str, grade: int, subject: str):
        """
        Background task to process uploaded curriculum PDF.
        This function runs asynchronously and won't block the API response.
        """
        print(f"--- Background task started: Processing PDF '{original_filename}'. ---", flush=True)
        pool_handler = None
        try:
            # --- Setup Paths ---
            project_root = Path(__file__).parent
            embeddings_dir = project_root / "embeddings"
            main_json_path = project_root / "All_Curriculums_grouped.json"
            
            embeddings_dir.mkdir(exist_ok=True)

            # --- Create Dependencies ---
            pool_handler = ConnectionPool(
                dbname=os.getenv("POSTGRES_DB"), 
                user=os.getenv("POSTGRES_USER"), 
                password=os.getenv("POSTGRES_PASSWORD"), 
                host=os.getenv("DB_HOST", "postgres"), 
                port=int(os.getenv("DB_PORT", 5432))
            )
            ingestion_service = DataIngestionService(pool_handler=pool_handler)
            
            # --- 1. Save and Run Pipeline ---
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf:
                temp_pdf.write(pdf_bytes)
                temp_pdf_path = temp_pdf.name
            
            print(f"--- Background task: Saved temp PDF to {temp_pdf_path} ---", flush=True)
            temp_csv_path, temp_json_path = run_processing_pipeline(temp_pdf_path, grade, subject)
            
            # --- 2. Save the generated CSV ---
            csv_filename = Path(temp_csv_path).name
            csv_dest_path = embeddings_dir / csv_filename
            shutil.move(temp_csv_path, csv_dest_path)
            print(f"--- Background task: Saved new embeddings to '{csv_dest_path}' ---", flush=True)

            # ==========================================================
            # === CORRECTED LOGIC: APPEND NEW CURRICULUM TO EXISTING ===
            # ==========================================================
            # --- 3. Read both JSON files ---
            print("--- Background task: Reading generated JSON structure... ---", flush=True)
            with open(temp_json_path, 'r', encoding='utf-8') as f:
                new_structure_data = json.load(f)
            
            print(f"--- Background task: New structure contains keys: {list(new_structure_data.keys())} ---", flush=True)
            
            # Load existing main JSON or start with empty dict
            try:
                with open(main_json_path, 'r', encoding='utf-8') as f:
                    existing_structure_data = json.load(f)
                print(f"--- Background task: Loaded existing structure with {len(existing_structure_data)} curricula ---", flush=True)
            except FileNotFoundError:
                print("--- Background task: Main JSON file not found. Creating new one. ---", flush=True)
                existing_structure_data = {}
            except json.JSONDecodeError:
                print("--- Background task: Main JSON file corrupted. Starting fresh. ---", flush=True)
                existing_structure_data = {}
            
            # Append new curriculum keys to the existing structure
            for curriculum_key, curriculum_content in new_structure_data.items():
                if curriculum_key in existing_structure_data:
                    print(f"--- WARNING: Key '{curriculum_key}' already exists. Overwriting. ---", flush=True)
                else:
                    print(f"--- Background task: Adding new curriculum '{curriculum_key}' to main JSON. ---", flush=True)
                
                existing_structure_data[curriculum_key] = curriculum_content
            
            # Write the updated data back to the file
            with open(main_json_path, 'w', encoding='utf-8') as f:
                json.dump(existing_structure_data, f, indent=2, ensure_ascii=False)
            
            print(f"--- Background task: Main JSON now contains {len(existing_structure_data)} curricula ---", flush=True)
            # ==========================================================

            # --- 4. Ingest structure into DB ---
            print("--- Background task: Ingesting new structure into DB... ---", flush=True)
            db_formatted_structure = convert_json_to_db_format(new_structure_data)
            ingestion_service.ingest_curriculum_structure(db_formatted_structure)

            # --- 5. Ingest embeddings into DB ---
            print("--- Background task: Ingesting new embeddings into DB... ---", flush=True)
            embeddings_df = pd.read_csv(csv_dest_path)
            ingestion_service.ingest_embeddings_from_csv(embeddings_df)

            print("--- Background task: Verifying database insertions... ---", flush=True)
            from services.pgvector_service import PGVectorService
            pgvector_service = PGVectorService(pool_handler)
            pgvector_service.verify_recent_insertions()
            
            # --- 6. Cleanup ---
            os.unlink(temp_pdf_path)
            os.unlink(temp_json_path)
            print("--- Background task: Cleaned up temporary files ---", flush=True)

            print("--- ✅ Background task completed successfully. ---", flush=True)

        except Exception as e:
            import traceback
            print(f"--- ❌ FATAL ERROR in background task: {e} ---", flush=True)
            print(f"--- Traceback: {traceback.format_exc()} ---", flush=True)
        finally:
            if pool_handler:
                pool_handler.close_all()
                print("--- Background task: Database connection pool closed. ---", flush=True)

    @app.on_event("startup")
    async def startup_event():
        # Access the container from app state to print config on startup
        container = app.state.container
        print("MinIO Endpoint:", container.config.minio_endpoint)
        print("MinIO Bucket:", container.config.minio_bucket)
        print("OpenAI Service Available:", container.openai_service.is_available())
        print("Agent Service Available:", container.agent_service.is_available())

    @app.get("/chat-interface")
    async def serve_audio_recorder():
        """Serve the audio recorder HTML file"""
        try:
            static_file = Path("static/audio-recorder.html")
            if static_file.exists():
                return FileResponse(static_file)
            current_file = Path("audio-recorder.html")
            if current_file.exists():
                return FileResponse(current_file)
            raise HTTPException(status_code=404, detail="Audio recorder interface not found")
        except Exception as e:
            print(f"Error serving audio recorder: {e}")
            raise HTTPException(status_code=500, detail=f"Error serving interface: {str(e)}")

    @app.get("/curriculum-upload")
    async def serve_curriculum_upload():
        """Serve the curriculum upload HTML file"""
        try:
            static_file = Path("static/curriculum_PDF_uploader.html")
            if static_file.exists():
                return FileResponse(static_file)
            current_file = Path("curriculum_PDF_uploader.html")
            if current_file.exists():   
                return FileResponse(current_file)
            raise HTTPException(status_code=404, detail="Curriculum upload interface not found")
        except Exception as e:
            print(f"Error serving curriculum upload interface: {e}")
            raise HTTPException(status_code=500, detail=f"Error serving interface: {str(e)}")
        

        
    @app.post("/chat")
    async def chat_handler(
        request: Request,
        file: Optional[UploadFile] = File(None), 
        text: Optional[str] = Form(None),
        student_id: str = Form("student_001"),
        game_context: Optional[str] = Form(None) 
    ):
        """Handles incoming chat messages using the shared container instance."""
        container = request.app.state.container
        try:
            if not student_id.strip():
                raise HTTPException(status_code=400, detail="Student ID is required")
            
            result = container.chat_service.process_message(
                student_id=student_id, 
                file=file, 
                text=text,
                game_context=game_context
            )
            return result
        except Exception as e:
            print(f"Error in chat handler: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Chat processing error: {str(e)}")

    @app.get("/get-audio-response")
    async def get_audio_response(request: Request, student_id: str = "student_001"):
        """Fetches the agent's text and audio response using the shared container."""
        container = request.app.state.container
        try:
            result = container.response_service.get_agent_response(student_id=student_id)
            if hasattr(result, 'status_code'):
                return result
            # This should be unreachable if response_service always returns a Response object
            return result
        except Exception as e:
            print(f"Error getting audio response: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Audio response error: {str(e)}")

    @app.post("/process-curriculum", status_code=202)
    async def process_curriculum_webhook(
        background_tasks: BackgroundTasks,
        grade: int = Form(...), 
        subject: str = Form(...),
        file: UploadFile = File(...)
                                    ):
        """
        Accepts a PDF and adds a background task to process it.
        Returns immediately.
        """
        pdf_bytes = await file.read()
        
        # --- THIS IS THE FIX ---
        # We now pass BOTH required arguments to the add_task method.
        background_tasks.add_task(
            process_pdf_curriculum_in_background, 
            pdf_bytes, 
            file.filename,
            grade,
            subject
        )
        
        # Return immediately to the user
        return {"status": "processing_started", "message": "The curriculum is being processed in the background."}


    @app.post("/mcq/generate")
    async def generate_mcqs_handler(
        request: Request,
        grade: int = Form(...),
        subject: str = Form(...),
        unit: str = Form(...),
        concept: str = Form(...),
        count: int = Form(5),
        is_arabic: bool = Form(False)
    ):
        """
        Generates and stores a new set of MCQs for a specific topic.
        """
        container = request.app.state.container
        try:
            generated_questions = container.agent_service.generate_and_store_mcqs(
                grade=grade,
                subject=subject,
                unit=unit,
                concept=concept,
                num_questions=count,
                is_arabic=is_arabic
            )
            return {
                "status": "success",
                "message": f"Successfully generated and stored {len(generated_questions)} MCQs.",
                "questions": generated_questions
            }
        except HTTPException as e:
            raise e # Re-raise FastAPI specific exceptions
        except Exception as e:
            logger.error(f"Error in generate_mcqs_handler: {e}")
            raise HTTPException(status_code=500, detail=str(e))

    @app.get("/mcq")
    async def get_mcqs_handler(
        request: Request,
        grade: int,
        subject: str,
        unit: str,
        concept: str,
        is_arabic: bool,
        # Make limit optional. If not provided, it will be None.
        limit: Optional[int] = None 
    ):
        """
        Retrieves existing MCQs for a specific topic and language from the database.
        If no limit is provided, retrieves all questions.
        """
        container = request.app.state.container
        try:
            questions = container.agent_service.pgvector.get_mcqs(
                grade=grade,
                subject=subject,
                unit=unit,
                concept=concept,
                is_arabic=is_arabic,
                limit=limit # Pass the limit (which could be None)
            )
            return {
                "status": "success",
                "count": len(questions),
                "questions": questions
            }
        except Exception as e:
            logger.error(f"Error in get_mcqs_handler: {e}")
            raise HTTPException(status_code=500, detail=str(e))
            
    @app.options("/get-audio-response")
    async def audio_response_options():
        """Handle preflight CORS requests for audio response endpoint"""
        return Response(status_code=204, headers={"Access-Control-Allow-Origin": "*", "Access-Control-Allow-Methods": "GET, OPTIONS", "Access-Control-Allow-Headers": "*", "Access-Control-Expose-Headers": "X-Response-Text"})

    @app.get("/health")
    async def health_check(request: Request):
        """Health check endpoint using the shared container."""
        container = request.app.state.container
        try:
            health_status = container.health_service.get_health_status()
            health_status.update({
                "openai_service_status": "available" if container.openai_service.is_available() else "unavailable",
                "agent_service_status": "available" if container.agent_service.is_available() else "unavailable",
                "minio_endpoint": container.config.minio_endpoint,
                "minio_bucket": container.config.minio_bucket
            })
            return health_status
        except Exception as e:
            print(f"Health check error: {e}")
            return {"status": "error", "message": str(e)}

    # Agent management endpoints
    @app.get("/conversation/stats")
    async def get_conversation_stats(request: Request, student_id: str = "student_001"):
        container = request.app.state.container
        try:
            return container.chat_service.get_agent_stats(student_id)
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))

    @app.post("/conversation/clear")
    async def clear_conversation(request: Request, student_id: str = Form("student_001")):
        container = request.app.state.container
        try:
            return container.chat_service.clear_conversation(student_id)
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))

    @app.post("/agent/system-prompt")
    async def set_system_prompt(req_body: dict, request: Request):
        container = request.app.state.container
        try:
            prompt = req_body.get("prompt", "")
            if not prompt:
                raise HTTPException(status_code=400, detail="System prompt cannot be empty")
            return container.chat_service.set_system_prompt(prompt)
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))

    @app.get("/agent/system-prompt")
    async def get_system_prompt(request: Request):
        container = request.app.state.container
        try:
            return {
                "system_prompt": container.agent_service.system_prompt,
                "status": "success"
            }
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))

    @app.get("/conversation/export")
    async def export_conversation(request: Request, student_id: str = "student_001"):
        container = request.app.state.container
        try:
            history = container.agent_service.export_conversation(student_id)
            return {
                "student_id": student_id,
                "messages": history,
                "total_messages": len(history)
            }
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))

    @app.post("/conversation/import")
    async def import_conversation(req_body: dict, request: Request):
        container = request.app.state.container
        try:
            student_id = req_body.get("student_id", "student_001")
            messages = req_body.get("messages", [])
            if not messages:
                raise HTTPException(status_code=400, detail="Messages list cannot be empty")
            container.agent_service.import_conversation(messages, student_id)
            return {"status": "success", "message": f"Imported {len(messages)} messages"}
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))

    @app.get("/debug/test-response")
    async def debug_test_response():
        """Debug endpoint to test response generation"""
        test_text = "This is a test response"
        encoded_text = base64.b64encode(test_text.encode('utf-8')).decode('utf-8')
        return Response(content=b"test audio data", media_type="audio/mpeg", headers={"X-Response-Text": encoded_text, "Access-Control-Expose-Headers": "X-Response-Text"})

    @app.get("/")
    async def root():
        """Root endpoint with API info"""
        return {"service": "Unified Chat API with Local Agent", "version": "2.2.0-lifespan", "status": "running"}
    
    return app

# Application entry point
app = create_app()

if __name__ == "__main__":
    # For development
    uvicorn.run(
        "main:app", 
        host="0.0.0.0", 
        port=int(os.environ.get("PORT", 8000)),
        reload=True
    )