from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from starlette.requests import Request
import os
import time
from pathlib import Path
from typing import Optional, List

# Import your backend modules
from backend.file_handler import save_upload
from backend.extractors import extract_text
from backend.qa_engine import QAEngine
from backend.image_processor import ImageProcessor
from backend.response_formatter import ResponseFormatter

app = FastAPI(
    title="Intelligent QA Service",
    description="Question answering for documents and images"
)



# Try to set Hugging Face token if available in environment
huggingface_token = os.environ.get("HF_TOKEN")
if huggingface_token:
    from huggingface_hub import login
    login(token=huggingface_token)

# Initialize models with fallback options
try:
    qa_engine = QAEngine(model_name="distilbert-base-cased-distilled-squad")  # Use a public model
except Exception as e:
    print(f"Error initializing QA engine: {str(e)}")
    # Fallback to a simpler implementation if needed
    from backend.qa_engine import SimpleQAEngine
    qa_engine = SimpleQAEngine()


try:
    image_processor = ImageProcessor()
except Exception as e:
    print(f"Error initializing Image Processor: {str(e)}")
    # Create a fallback image processor if needed
    from backend.image_processor import SimpleImageProcessor
    image_processor = SimpleImageProcessor()



formatter = ResponseFormatter()



# Mount static files and templates
templates = Jinja2Templates(directory="frontend/templates")
app.mount("/static", StaticFiles(directory="frontend/static"), name="static")

@app.get("/")
async def read_root(request: Request):
    """Render the main page"""
    return templates.TemplateResponse("index.html", {"request": request})

@app.post("/api/document-qa")
async def document_qa(
    file: UploadFile = File(...),
    question: str = Form(...)
):
    """Process document and answer question"""
    try:
        # Save the uploaded file
        file_id, file_name = save_upload(file)
        file_path = Path(f"/tmp/uploads/{file_name}")
        
        # Extract text from document
        document_text = extract_text(str(file_path))
        
        # Get answer from QA engine
        if isinstance(document_text, dict):
            # Handle structured document text
            # This is a simplistic approach - you'd need to convert the 
            # structured content to plain text for the QA engine
            if "content" in document_text:
                if isinstance(document_text["content"], list):
                    if isinstance(document_text["content"][0], dict):
                        # Handle docx structure
                        text = " ".join([p["text"] for p in document_text["content"]])
                    else:
                        # Handle txt structure
                        text = " ".join(document_text["content"])
                else:
                    text = str(document_text["content"])
            else:
                text = str(document_text)
        else:
            # Plain text from PDF or PPTX
            text = document_text
            
        qa_result = qa_engine.answer_question(text, question)
        qa_result["timestamp"] = time.time()
        
        # Format response
        response = formatter.format_document_qa_response(qa_result, file.filename)
        
        return JSONResponse(content=response)
        
    except Exception as e:
        error_response = formatter.format_error_response(str(e))
        return JSONResponse(content=error_response, status_code=error_response["status_code"])

@app.post("/api/image-qa")
async def image_qa(
    file: UploadFile = File(...),
    question: str = Form(...)
):
    """Process image and answer question"""
    try:
        print(f"Received image: {file.filename}, size: {file.size}, question: {question}")
        
        # Validate file is an image
        if not file.content_type.startswith('image/'):
            print(f"Invalid content type: {file.content_type}")
            return JSONResponse(
                content={"error": "File must be an image", "status_code": 400},
                status_code=400
            )
        
        # Save the uploaded file
        file_id, file_name = save_upload(file)
        file_path = Path(f"/tmp/uploads/{file_name}")
        print(f"Saved image to: {file_path}")
        
        if not file_path.exists():
            print(f"File not saved properly at {file_path}")
            return JSONResponse(
                content={"error": "File could not be saved", "status_code": 500},
                status_code=500
            )
        
        # Process the image
        vqa_result = image_processor.answer_image_question(str(file_path), question)
        vqa_result["timestamp"] = time.time()
        
        # Format response
        response = formatter.format_image_qa_response(vqa_result, file.filename)
        
        return JSONResponse(content=response)
        
    except Exception as e:
        import traceback
        print(f"Error in image_qa: {str(e)}")
        print(traceback.format_exc())
        error_response = formatter.format_error_response(str(e))
        return JSONResponse(content=error_response, status_code=error_response.get("status_code", 500))

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