app.py

from __future__ import annotations
import os, io, re, json, time, mimetypes, tempfile, string
from typing import List, Union, Tuple, Any, Iterable

from PIL import Image
import pandas as pd
import gradio as gr
import google.generativeai as genai
import requests
import pdfplumber

# ================== CONFIG ==================
DEFAULT_API_KEY = "AIzaSyBbK-1P3JD6HPyE3QLhkOps6_-Xo3wUFbs"

INTERNAL_MODEL_MAP = {
    "Gemini 2.5 Flash": "gemini-2.5-flash",
    "Gemini 2.5 Pro":   "gemini-2.5-pro",
}
EXTERNAL_MODEL_NAME = "prithivMLmods/Camel-Doc-OCR-062825 (External)"

try:
    RESAMPLE = Image.Resampling.LANCZOS
except AttributeError:
    RESAMPLE = Image.LANCZOS

PROMPT_FREIGHT_JSON = """
Please analyze the freight rate table in the file I provide and convert it into JSON in the following structure:
{
  "shipping_line": "...",
  "shipping_line_code": "...",
  "shipping_line_reason": "Why this carrier is chosen?",
  "fee_type": "Air Freight",
  "valid_from": ...,
  "valid_to": ...,
  "charges": [
    {
      "frequency": "...",
      "package_type": "...",
      "aircraft_type": "...",
      "direction": "Export or Import or null",
      "origin": "...",
      "destination": "...",
      "charge_name": "...",
      "charge_code": "...",
      "charge_code_reason": "...",
      "cargo_type": "...",
      "currency": "...",
      "transit": "...",
      "transit_time": "...",
      "weight_breaks": {
        "M": ...,
        "N": ...,
        "+45kg": ...,
        "+100kg": ...,
        "+300kg": ...,
        "+500kg": ...,
        "+1000kg": ...,
        "other": {
          key: value
        },
        "weight_breaks_reason":"Why chosen weight_breaks?"
      },
      "remark": "..."
    }
  ],
  "local_charges": [
    {
      "charge_name": "...",
      "charge_code": "...",
      "unit": "...",
      "amount": ...,
      "remark": "..."
    }
  ]
}
### Date rules
- valid_from format:
  - `DD/MM/YYYY` (if full date)
  - `01/MM/YYYY` (if month+year only)
  - `01/01/YYYY` (if year only)
  - `UFN` if missing
- valid_to:
  - exact `DD/MM/YYYY` if present
  - else `UFN`
STRICT RULES:
- ONLY return a single JSON object as specified above.
- All rates must exactly match the corresponding weight break columns (M,N,45kg, 100kg, 300kg, 500kg, 1000kg, etc.). set null if N/A. No assumptions or interpolations.
- If the table shows "RQ" or similar, set value as "RQST".
- Group same-price destinations into one record separated by "/".
- Always use IATA code for origin and destination.
- Flight number (e.g. ZH118) is not charge code.
- Frequency: D[1-7]; 'Daily' = D1234567. Join multiple (e.g. D3,D4→D34).
- If local charges exist, list them.
- If validity missing, set null.
- Direction: Export if origin is Vietnam (SGN, HAN, DAD...), else Import.
- Provide short plain English reasons for "shipping_line_reason" & "charge_code_reason".
- Replace commas in remarks with semicolons.
- Only return JSON.
"""

# ================== HELPERS ==================
import fitz  # PyMuPDF

def pdf_to_images(pdf_bytes: bytes) -> list[Image.Image]:
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    pages = []
    for p in doc:
        pix = p.get_pixmap(dpi=200)
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        pages.append(img)
    return pages

def ensure_rgb(im: Image.Image) -> Image.Image:
    return im.convert("RGB") if im.mode != "RGB" else im

def _read_file_bytes(upload: Union[str, os.PathLike, dict, object] | None) -> bytes:
    if upload is None:
        raise ValueError("No file uploaded.")
    if isinstance(upload, (str, os.PathLike)):
        with open(upload, "rb") as f:
            return f.read()
    if isinstance(upload, dict) and "path" in upload:
        with open(upload["path"], "rb") as f:
            return f.read()
    if hasattr(upload, "read"):
        return upload.read()
    raise TypeError(f"Unsupported file object: {type(upload)}")

def _guess_name_and_mime(file, file_bytes: bytes) -> Tuple[str, str]:
    if isinstance(file, (str, os.PathLike)):
        filename = os.path.basename(str(file))
    elif isinstance(file, dict) and "name" in file:
        filename = os.path.basename(file["name"])
    elif isinstance(file, dict) and "path" in file:
        filename = os.path.basename(file["path"])
    else:
        filename = "upload.bin"
    mime, _ = mimetypes.guess_type(filename)
    if not mime:
        if len(file_bytes) >= 4 and file_bytes[:4] == b"%PDF":
            mime = "application/pdf"
            if not filename.lower().endswith(".pdf"):
                filename += ".pdf"
        else:
            mime = "image/png"
    return filename, mime

# ================== PDF CHECK STEP ==================
def check_pdf_structure(file_bytes: bytes) -> str:
    """Kiểm tra nhanh file PDF có phải bảng nhiều cột, nhiều trang không."""
    try:
        with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
            if len(pdf.pages) <= 2:
                return "không"
            table_pages = 0
            for page in pdf.pages[:3]:
                tables = page.find_tables()
                if tables and len(tables) > 0:
                    table_pages += 1
            if table_pages >= 1:
                return "có"
            text = "\n".join([(p.extract_text() or "") for p in pdf.pages[:2]])
            num_tokens = sum(ch.isdigit() for ch in text)
            line_count = len(text.splitlines())
            if num_tokens > 100 and line_count > 20:
                return "có"
        return "không"
    except Exception as e:
        print("PDF check error:", e)
        return "không"

# ================== OCR CORE (Gemini) ==================
def run_process_internal_base_v2(file_bytes, filename, mime, question, model_choice, temperature, top_p, batch_size=3):
    api_key = os.environ.get("GOOGLE_API_KEY", DEFAULT_API_KEY)
    if not api_key:
        return "ERROR: Missing GOOGLE_API_KEY.", None
    genai.configure(api_key=api_key)
    model_name = INTERNAL_MODEL_MAP.get(model_choice, "gemini-2.5-flash")
    model = genai.GenerativeModel(model_name=model_name,
                                  generation_config={"temperature": float(temperature), "top_p": float(top_p)})

    if file_bytes[:4] == b"%PDF":
        pages = pdf_to_images(file_bytes)
    else:
        pages = [Image.open(io.BytesIO(file_bytes))]

    user_prompt = (question or "").strip() or PROMPT_FREIGHT_JSON
    all_json_results, all_text_results = [], []
    previous_header_json = None

    def _safe_text(resp):
        try:
            return resp.text
        except:
            return ""

    for i in range(0, len(pages), batch_size):
        batch = pages[i:i+batch_size]
        uploaded = []
        for im in batch:
            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
                im.save(tmp.name)
                up = genai.upload_file(path=tmp.name, mime_type="image/png")
                up = genai.get_file(up.name)
                uploaded.append(up)

        context_prompt = user_prompt
        resp = model.generate_content([context_prompt] + uploaded)
        text = _safe_text(resp)
        all_text_results.append(text)
        for up in uploaded:
            try:
                genai.delete_file(up.name)
            except:
                pass

    return "\n\n".join(all_text_results), None

# ================== EXTERNAL API (nếu có) ==================
def run_process_external(file_bytes, filename, mime, question, api_url, temperature, top_p):
    if not api_url:
        return "ERROR: Missing external API endpoint.", None
    data = {"prompt": question or "", "temperature": str(temperature), "top_p": str(top_p)}
    files = {"file": (filename, file_bytes, mime)}
    r = requests.post(api_url, files=files, data=data, timeout=60)
    if r.status_code >= 400:
        return f"ERROR: External API HTTP {r.status_code}: {r.text[:200]}", None
    return r.text, None

# ================== MAIN ROUTER (đã thêm STEP CHECK) ==================
def run_process(file, question, model_choice, temperature, top_p, external_api_url):
    """
    Router (có bước kiểm tra PDF/table trước khi xử lý):
      - Nếu PDF nhiều trang/nhiều bảng -> extract trước (pdfplumber)
      - Ngược lại -> OCR trực tiếp Gemini
    """
    try:
        if file is None:
            return "ERROR: No file uploaded.", None

        file_bytes = _read_file_bytes(file)
        filename, mime = _guess_name_and_mime(file, file_bytes)

        # STEP 1️⃣: Check PDF structure
        if mime == "application/pdf" or file_bytes[:4] == b"%PDF":
            check_result = check_pdf_structure(file_bytes)
            print(f"[PDF Check] {filename}: {check_result}")

            if check_result == "có":
                try:
                    print("➡️ PDF có nhiều cột/nhiều trang → dùng pdfplumber extract trước rồi Gemini.")
                    all_dfs = []
                    saved_header = None
            
                    with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
                        for page_idx, page in enumerate(pdf.pages, start=1):
                            print(f"📄 Đang xử lý trang {page_idx}...")
            
                            table = page.extract_table({
                                "vertical_strategy": "lines",
                                "horizontal_strategy": "text",
                                "snap_tolerance": 3,
                                "intersection_tolerance": 5,
                            })
            
                            if not table or len(table) < 2:
                                print(f"⚠️ Trang {page_idx}: Không phát hiện bảng hợp lệ.")
                                continue
            
                            header = table[0]
                            rows = table[1:]
            
                            # Lưu header đầu tiên
                            if saved_header is None:
                                saved_header = header
                                print(f"✅ Trang {page_idx}: Lưu header đầu tiên: {saved_header}")
            
                            # Nếu trang sau không có header rõ → dùng header cũ
                            if len(header) < len(saved_header) or "REGION" not in header[0]:
                                print(f"↩️ Trang {page_idx}: Không có header rõ ràng, dùng lại header trước.")
                                header = saved_header
                                rows = table
                            else:
                                saved_header = header  # cập nhật header hợp lệ
            
                            if len(rows) == 0:
                                print(f"⚠️ Trang {page_idx}: Không có dữ liệu dưới header.")
                                continue
            
                            try:
                                df = pd.DataFrame(rows, columns=header)
                                all_dfs.append(df)
                                print(f"✅ Trang {page_idx}: {len(df)} dòng được thêm.")
                            except Exception as e:
                                print(f"❌ Lỗi tạo DataFrame ở trang {page_idx}: {e}")
            
                    if all_dfs:
                        final_df = pd.concat(all_dfs, ignore_index=True).dropna(how="all").reset_index(drop=True)
                        print(f"✅ Tổng cộng {len(final_df)} dòng được trích xuất từ PDF.")
            
                        # Xuất ra file tạm (Excel + JSON)
                        base_name = os.path.splitext(filename)[0]
                        tmp_dir = tempfile.gettempdir()
                        # json_path = os.path.join(tmp_dir, f"{base_name}.json")
                        # excel_path = os.path.join(tmp_dir, f"{base_name}.xlsx")
            
                        # final_df.to_json(json_path, orient="records", force_ascii=False, indent=2)
                        # final_df.to_excel(excel_path, index=False)
            
                        # print(f"✅ Xuất JSON:  {json_path}")
                        # print(f"✅ Xuất Excel: {excel_path}")
            
                        # Convert bảng thành CSV text để Gemini đọc tiếp
                        table_text = final_df.to_csv(index=False)
                        print(f"✅ Đang Gen text từ file CSV")
                        question = (
                            f"{PROMPT_FREIGHT_JSON}\n"
                            "Below is the table text extracted from the PDF (CSV format):\n"
                            f"{table_text}\n\n"
                            "Please convert this into valid JSON as per the schema."
                        )
                    else:
                        print("⚠️ Không có bảng hợp lệ để extract bằng pdfplumber.")
            
                except Exception as e:
                    print("❌ pdfplumber extract failed:", e)


        # STEP 2️⃣: Route model
        if model_choice == EXTERNAL_MODEL_NAME:
            return run_process_external(
                file_bytes=file_bytes, filename=filename, mime=mime,
                question=question, api_url=external_api_url,
                temperature=temperature, top_p=top_p
            )

        return run_process_internal_base_v2(
            file_bytes=file_bytes, filename=filename, mime=mime,
            question=question, model_choice=model_choice,
            temperature=temperature, top_p=top_p
        )

    except Exception as e:
        return f"ERROR: {type(e).__name__}: {str(e)}", None

# ================== UI ==================
def main():
    with gr.Blocks(title="OCR Multi-Agent System") as demo:
        file = gr.File(label="Upload PDF/Image")
        question = gr.Textbox(label="Prompt", lines=2)
        model_choice = gr.Dropdown(choices=[*INTERNAL_MODEL_MAP.keys(), EXTERNAL_MODEL_NAME],
                                   value="Gemini 2.5 Flash", label="Model")
        temperature = gr.Slider(0.0, 2.0, value=0.2, step=0.05)
        top_p = gr.Slider(0.0, 1.0, value=0.95, step=0.01)
        external_api_url = gr.Textbox(label="External API URL", visible=False)
        output_text = gr.Code(label="Output", language="json")
        run_btn = gr.Button("🚀 Process")

        run_btn.click(
            run_process,
            inputs=[file, question, model_choice, temperature, top_p, external_api_url],
            outputs=[output_text, gr.State()]
        )

    return demo

demo = main()

if __name__ == "__main__":
    demo.launch()