import gradio as gr
from fastapi import FastAPI, UploadFile, File, Body
from fastapi.middleware.cors import CORSMiddleware
import pandas as pd
import os
import uvicorn
import json

from ocr import scan_receipt
from predict import predict_expense
from behavior import analyze_behavior
from chat import chat_response


# ---------------- INSTALL TESSERACT ----------------
if not os.path.exists("/usr/bin/tesseract"):
    os.system("apt-get update && apt-get install -y tesseract-ocr")


# ---------------- FASTAPI ----------------
api = FastAPI()

api.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# 🔹 OCR API
@api.post("/scan")
async def scan(file: UploadFile = File(...)):
    content = await file.read()
    return scan_receipt(content)


# 🔹 Prediction API
@api.post("/predict")
async def predict(data: list = Body(...)):
    df = pd.DataFrame(data)
    return predict_expense(df)


# 🔹 Behavior API
@api.post("/behavior")
async def behavior(data: list = Body(...)):
    df = pd.DataFrame(data)
    return analyze_behavior(df)


# 🔹 Chat API
@api.post("/chat")
async def chat(req: dict):
    query = req.get("query")
    token = req.get("token")

    if not token:
        return {"error": "Missing token"}

    response = chat_response(query, token)
    return {"response": response}


# ---------------- GRADIO UI ----------------
def ocr_ui(file):
    return scan_receipt(open(file.name, "rb").read())


def predict_ui(data):
    df = pd.DataFrame(json.loads(data))
    return predict_expense(df)


def behavior_ui(data):
    df = pd.DataFrame(json.loads(data))
    return analyze_behavior(df)


def chat_ui(query, token):
    if not token or token.strip() == "":
        return "❌ Please provide a valid access token."

    return chat_response(query, token)


with gr.Blocks() as ui:
    gr.Markdown("# 💰 Expense AI")

    # OCR
    with gr.Tab("OCR"):
        image = gr.File(label="Upload Receipt")
        output = gr.JSON(label="OCR Result")
        gr.Button("Scan").click(ocr_ui, inputs=image, outputs=output)

    # Prediction
    with gr.Tab("Prediction"):
        inp = gr.Textbox(label="Enter JSON data")
        out = gr.JSON(label="Prediction Result")
        gr.Button("Predict").click(predict_ui, inputs=inp, outputs=out)

    # Behavior
    with gr.Tab("Behavior"):
        inp2 = gr.Textbox(label="Enter JSON data")
        out2 = gr.JSON(label="Behavior Analysis")
        gr.Button("Analyze").click(behavior_ui, inputs=inp2, outputs=out2)

    # Chat
    with gr.Tab("Chat"):
        gr.Markdown("### 🔐 Enter your Supabase access token")
        token_input = gr.Textbox(label="Access Token", type="password")
        chat_in = gr.Textbox(label="Ask your financial question")
        chat_out = gr.Textbox(label="AI Response")

        gr.Button("Ask AI").click(chat_ui, inputs=[chat_in, token_input], outputs=chat_out)


# ---------------- COMBINE ----------------
app = gr.mount_gradio_app(api, ui, path="/")


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