from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from mpxpy.mathpix_client import MathpixClient
import os
import shutil
from dotenv import load_dotenv

# Load .env file
load_dotenv()

app = FastAPI()
# Initialize Mathpix client (will use ~/.mpx/config or env variables)
client = MathpixClient()

def detect_image_type(result: dict) -> str:
    """
    Detect whether Mathpix response is LaTeX/math or plain text using
    a robust logic based on API response fields.
    """
    # 1. Prioritize the top-level content_type field if available.
    #    This is the most direct and reliable classification from Mathpix.
    content_type = result.get("content_type")
    if content_type == "math":
        return "latex"
    elif content_type == "text":
        return "text"

    # 2. If content_type is not conclusive, analyze the line-by-line data
    #    with a majority rule.
    line_data = result.get("line_data", [])
    if line_data:
        math_lines = 0
        text_lines = 0
        for line in line_data:
            line_type = line.get("type")
            if line_type == "math":
                math_lines += 1
            elif line_type == "text":
                text_lines += 1
        
        total_lines = math_lines + text_lines
        if total_lines > 0:
            math_ratio = math_lines / total_lines
            # If a significant portion is math, classify as latex. This threshold can be adjusted.
            if math_ratio > 0.4:
                return "latex"
            else:
                return "text"

    # 3. As a final fallback, use the top-level confidence scores.
    #    Note: These are only available for certain formats.
    latex_confidence = result.get("latex_confidence", 0)
    text_confidence = result.get("text_confidence", 0)
    if latex_confidence > text_confidence:
        return "latex"
    
    # 4. Default to 'text' if all other checks are inconclusive.
    return "text"


# --- API Endpoint ---
@app.post("/image-to-latex")
async def image_to_latex(file: UploadFile = File(...)):
    try:
        UPLOAD_DIR = "received_images"  # your folder name

        # make sure the folder exists
        os.makedirs(UPLOAD_DIR, exist_ok=True)

        # Save file in folder
        save_path = os.path.join(UPLOAD_DIR, file.filename)

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

        # Send to Mathpix with better options
        image = client.image_new(
            file_path=save_path,
            formats=["text", "latex_styled", "data", "mathml"],
            data_options={"include_asciimath": True, "include_latex": True, "include_mathml": True, "include_table_html": True},
            include_line_data=True,
        )

        # return image

        result = image.results()
        # return result
        # Detect type (latex or text)
        img_type = detect_image_type(result)
        
        print("img_type", img_type)
        print("text", result.get("text", ""))
        print("latex", result.get("latex_styled", ""))

        return JSONResponse(
            content={
                "status": True,
                "image_type": img_type,
                "latex": result.get("latex_styled", ""),
                "text": result.get("text", ""),
                "asciimath": result.get("asciimath", ""),
                "mathml": result.get("mathml", ""),
                "raw_result": result,
                "line_data": result.get("line_data", []),
            }
        )

    except Exception as e:
        print("Mathpix API error:", str(e))
        return JSONResponse(
            content={"status": False, "message": str(e)}, status_code=500
        )