import base64
from pathlib import Path

from fastapi import APIRouter, UploadFile, File, status

from app.schemas import ResponseWrapper
from app.exceptions import HTTPException
from app.configs.settings import settings
from app.schemas.ocr import OCRResponse, ImageOCRSchema


from langchain_openai.chat_models.base import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate

router = APIRouter()

ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif", ".tiff", ".tif"}
ALLOWED_CONTENT_TYPES = {
    "image/jpeg",
    "image/png",
    "image/webp",
    "image/bmp",
    "image/gif",
    "image/tiff",
    "image/tif",
}

async def get_chat_prompt(base64: str) -> ChatPromptTemplate:

    input_content = [
        {
            "type": "image_url",
            "image_url": {
            "url": base64
            }
        }
    ]

    return ChatPromptTemplate.from_messages([
        SystemMessagePromptTemplate.from_template("You are a helpful assistant that can extract information from images or documents. Try to avoid short form filling and try to fill full form if possible else fill the short form"),
        HumanMessagePromptTemplate.from_template(input_content)
    ])

@router.post(
    "/ocr",
    tags=["OCR"],
    response_model=ResponseWrapper[OCRResponse],
    summary="OCR document.",
    status_code=200,
)
async def ocr(file: UploadFile = File(...)) -> ResponseWrapper[OCRResponse]:

    # Validate file extension
    extension = Path(file.filename or "").suffix.lower()
    if extension not in ALLOWED_EXTENSIONS:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Only image files (.jpg, .jpeg, .png, .webp, .bmp, .gif, .tiff, .tif) are allowed.",
        )

    # Validate content type
    if file.content_type not in ALLOWED_CONTENT_TYPES:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Invalid file type. Please upload a valid image.",
        )

    # Read uploaded image
    image_bytes = await file.read()

    if not image_bytes:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Uploaded file is empty.",
        )

    # Convert image to Base64 Data URL
    base64_image = (
        f"data:{file.content_type};base64,"
        f"{base64.b64encode(image_bytes).decode('utf-8')}"
    )

    llm = ChatOpenAI(api_key=settings.OPENAI_API_KEY, model="gpt-5.4")

    structured_llm = llm.with_structured_output(ImageOCRSchema)

    chain = await get_chat_prompt(base64_image) | structured_llm

    result: ImageOCRSchema = await chain.ainvoke({})

    return ResponseWrapper(
        status=True, 
        message="SUCCESS", 
        data=OCRResponse(
            firstName=result.first_name, 
            lastName=result.last_name, 
            email=result.email, 
            phoneNumber=result.phone, 
            address=result.address, 
            city=result.city, 
            state=result.state, 
            zipCode=result.zip, 
            country=result.country, 
            documentType=result.document_type, 
            documentNumber=result.document_number, 
            dateOfBirth=result.dob
        )
    )