from fastapi import FastAPI, Request, UploadFile, File, Form
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, JSONResponse
import shutil
import os
import tempfile
import asyncio
from typing import List, Optional

from LM import run_pdf, run_images


try:
    with open("InvoicePrompt", 'r', encoding='utf-8') as f:
        SYS_INVOICE_PROMPT = f.read()
except FileNotFoundError:
    SYS_INVOICE_PROMPT = "请识别发票内容并返回JSON格式数据。"
    print("Warning: InvoicePrompt file not found.")

try:
    with open("OCR_MD", 'r', encoding='utf-8') as f:
        SYS_OCR_PROMPT = f.read()
except FileNotFoundError:
    SYS_OCR_PROMPT = "请识别图片中的文字并转换为Markdown格式。"
    print("Warning: OCR_MD file not found.")

app = FastAPI()

templates = Jinja2Templates(directory="templates")

app.mount("/static", StaticFiles(directory="static"), name="static")


@app.get("/", response_class=HTMLResponse)
async def read_root(request: Request):

    return templates.TemplateResponse("index2.html", {"request": request})


@app.post("/api/recognize")
async def recognize_files(
        files: List[UploadFile] = File(...),
        mode: str = Form(...),
        custom_prompt: str = Form(None)
):
    results = []

    final_prompt = ""
    if mode == 'invoice':
        final_prompt = SYS_INVOICE_PROMPT
        print("Using System Invoice Prompt")
    elif mode == 'ocr':
        final_prompt = SYS_OCR_PROMPT
        print("Using System OCR Prompt")
    elif mode == 'custom':
        final_prompt = custom_prompt if custom_prompt else ""
        print(f"Using Custom Prompt (Length: {len(final_prompt)})")
    else:
        final_prompt = SYS_INVOICE_PROMPT

    with tempfile.TemporaryDirectory() as temp_dir:
        tasks = []

        for file in files:
            file_path = os.path.join(temp_dir, file.filename)

            with open(file_path, "wb") as buffer:
                shutil.copyfileobj(file.file, buffer)

            original_name = file.filename
            ext = os.path.splitext(original_name)[1].lower()

            if ext == ".pdf":
                tasks.append(run_pdf(original_name, file_path, prompt=final_prompt))

            elif ext in [".jpg", ".jpeg", ".png", ".bmp", ".webp"]:
                tasks.append(run_images([original_name], [file_path], prompt=final_prompt))

            else:
                results.append({
                    "image_name": original_name,
                    "result": "Unsupported file type",
                    "state": "error"
                })

        if tasks:
            processed_results = await asyncio.gather(*tasks, return_exceptions=True)

            for res in processed_results:
                if isinstance(res, Exception):
                    print(f"Internal Error: {str(res)}")
                    results.append({
                        "state": "error",
                        "result": f"Internal execution error: {str(res)}"
                    })
                else:
                    results.append(res)

    return JSONResponse(content={"data": results})


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