app.py

from __future__ import annotations

import os, io, re, json, time, mimetypes, tempfile
from typing import List, Union, Tuple
from PIL import Image
import pandas as pd
import gradio as gr
import google.generativeai as genai
#import requests
import pdfplumber
from pdf2image import convert_from_path
#import pytesseract
from concurrent.futures import ThreadPoolExecutor, as_completed
import fitz  # PyMuPDF
import multiprocessing

num_cpus = multiprocessing.cpu_count()
# ================== CONFIG ==================
DEFAULT_API_KEY = [
    "AIzaSyD0qjaoOJwrLeOz9Ko8Bi9vRgTy3AefTC8",
    # "AIzaSyAq7Wsi6fR0oWrJQbFkgGNdvxJTn8hWEzQ",
    # "AIzaSyDRWRwwnYJktCULH8d26mzD1Lv4l0CdQws",
    # "AIzaSyDW-x3kTWC7s2NJBOFDU7uC0vhKnREbANw",
    # "AIzaSyAq7Wsi6fR0oWrJQbFkgGNdvxJTn8hWEzQ",
    # "AIzaSyD0qjaoOJwrLeOz9Ko8Bi9vRgTy3AefTC8"
]
key_index = 0

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)"

PROMPT_FREIGHT_HEADER_JSON = """Vui lòng trích xuất tất cả thông tin tiêu đề (header), và ghi chú bên ngoài bảng giá trong tài liệu. 
Trả lời bằng tiếng Việt, ngắn gọn, rõ ràng và trình bày theo dạng danh sách. 
Đặc biệt, cần xác định và chuẩn hóa ngày hiệu lực (valid from / to) theo văn bản trong tài liệu, tuân thủ chính xác các quy tắc định dạng ngày như sau: 
 - DD/MM/YYYY, 01/MM/YYYY, 01/01/YYYY hoặc UFN nếu không có thông tin rõ ràng.
 - Nếu chi tiết các hãng không có ngày hiệu lực sẽ lấy thông tin trên header
 """
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": "...",
      "base_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": "...",
      "additional_cost": ...,
      "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 DATA RULES
============================================================

- ONLY return a single JSON object.
- All rates must match the weight break columns (M, N, +45kg, etc.).
- Use `null` if value is missing.
- "RQ" or similar → set as `"RQST"`.
- Group destinations with same rate using "/".
- Use IATA codes for `origin` and `destination`.
- Ignore flight numbers like "ZH118" for charge_code.
- Frequency format:
  - D[1-7] (e.g. D1, D2345, D1234567)
- Local charges: must include if found.
- Validity fields (`valid_from`, `valid_to`): use rules above.
- Direction: Export if from Vietnam (SGN, HAN, DAD...), otherwise Import.
- Provide plain English for `shipping_line_reason` and `charge_code_reason`.
- Replace commas in remarks with semicolons.
- RETURN ONLY JSON — no explanations.

============================================================
### PACKAGE TYPE & SURCHARGE LOGIC
============================================================

- Always treat **Carton** as the base rate.
- Generate derived **Pallet** (or SKID) surcharges if found in remarks/notes.

▶️ Rules:

1️⃣ **SKID shipment surcharge**  
If remark says:  
"SKID shipment: add 10 cents (apply for GEN & PER)"  
→ Add surcharge line (+0.10 USD/kg) for Pallet GEN/PER.

- Increase all weight breaks by that value.
- Keep origin, destination, etc. unchanged.
- Mention derivation in `remark`.

2️⃣ **Regional surcharge**  
E.g.:  
"For SKID shipment: EU +USD0.30/kg and rest +USD0.20/kg (exclude RGN, MAA)"  
→ Generate 2 surcharge lines accordingly.

3️⃣ **Carton = Pallet**  
If remark says:  
"Carton = Pallet"  
→ Copy Carton rates into Pallet.  
Set `additional_cost` = 0.

4️⃣ **As per remark**  
If remark says:  
"For specific route with package type: as per remark"  
→ Parse to determine logic.

============================================================
### DERIVED CHARGE GENERATION
============================================================

- Derived charges must be appended to `"charges"` array.
- Must include:
  - `"package_type": "Pallet"`
  - `"base_package_type": "Carton"`
  - `"additional_cost"` = numeric surcharge
  - `"remark"` stating derivation
- Other fields (origin, destination...) must match base record.
- DO NOT remove the Carton base record.

============================================================
### EXAMPLES
============================================================

Base:
{
  "package_type": "Carton",
  "cargo_type": "GEN",
  "origin": "SGN",
  "destination": "NRT",
  "currency": "USD",
  "weight_breaks": {
    "+45kg": 6.05,
    "+100kg": 5.30,
    "+300kg": 4.80
  },
  "remark": "Carton base rate"
}

Derived (from SKID remark):
{
  "package_type": "Pallet",
  "base_package_type": "Carton",
  "cargo_type": "GEN, PER",
  "currency": "USD",
  "origin": "SGN",
  "destination": "NRT",
  "additional_cost": 0.10,
  "weight_breaks": {
    "+45kg": 6.15,
    "+100kg": 5.40,
    "+300kg": 4.90
  },
  "remark": "Derived from Carton; SKID shipment: add 10 cents (apply for GEN & PER)"
}
"""

# ================== HELPERS ==================
def get_next_key():
    global key_index
    key = DEFAULT_API_KEY[key_index % len(DEFAULT_API_KEY)]
    key_index += 1
    
    return key

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

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]:
    filename = "upload.bin"
    if isinstance(file, (str, os.PathLike)):
        filename = os.path.basename(str(file))
    elif isinstance(file, dict):
        filename = os.path.basename(file.get("name") or file.get("path", filename))
    mime, _ = mimetypes.guess_type(filename)
    if not mime and file_bytes[:4] == b"%PDF":
        mime = "application/pdf"
        if not filename.lower().endswith(".pdf"):
            filename += ".pdf"
    return filename, mime or "application/octet-stream"

def safe_parse_json(text: str):
    cleaned = re.sub(r"```json|```", "", text).strip()
    try:
        return json.loads(cleaned)
    except json.JSONDecodeError as e:
        print(f"❌ Failed to parse JSON: {e}")
        print("📄 Raw text:\n", cleaned[:300])
        return None
def check_pdf_structure(file_bytes: bytes) -> str:
    """
    Phân tích PDF xem thuộc loại:
    - 0: "1_trang_1_hang"
    - 1: "nhieu_trang_1_hang"
    - 2: "nhieu_hang"
    - "khong_xac_dinh": nếu có lỗi
    """
    try:
        airline_pattern = re.compile(r"(.*?CARGO.*?RATE\s+EX\s+[A-Z]{3})", re.IGNORECASE)
        airline_headers = set()

        with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
            for page in pdf.pages:
                text = page.extract_text()
                if not text:
                    continue

                for line in text.splitlines():
                    match = airline_pattern.search(line.strip())
                    if match:
                        airline_name = match.group(1).strip().upper()
                        airline_headers.add(airline_name)

            total_pages = len(pdf.pages)

        if len(airline_headers) > 1:
            return 2
        elif total_pages > 1:
            return 1
        else:
            return 0

    except Exception as e:
        print(f"❌ Lỗi phân tích PDF: {e}")
        return "khong_xac_dinh"

# ================== PDF CHECK & SPLIT ==================
def split_excel_by_airline_header(excel_path, sheet_name=0):
    df = pd.read_excel(excel_path, header=None, sheet_name=sheet_name)
    airline_chunks = {}
    pattern = re.compile(r".*CARGO.*RATE EX HAN", re.IGNORECASE)
    start_indices, airline_names = [], []
    for i, row in df.iterrows():
        line = " ".join([str(cell) for cell in row if pd.notnull(cell)])
        if pattern.match(line):
            start_indices.append(i)
            airline_names.append(line.strip())
    start_indices.append(len(df))
    for i in range(len(airline_names)):
        chunk_df = df.iloc[start_indices[i]:start_indices[i+1]].reset_index(drop=True)
        airline_chunks[airline_names[i]] = chunk_df
    return airline_chunks

def export_pdf_to_excel(pdf_path: str, excel_output_path: str):
    all_data = []
    with pdfplumber.open(pdf_path) as pdf:
        for page_num, page in enumerate(pdf.pages, start=1):
            tables = page.extract_tables()
            for table in tables:
                df = pd.DataFrame(table)
                df["__page__"] = page_num
                all_data.append(df)
    if all_data:
        final_df = pd.concat(all_data, ignore_index=True)
        final_df.to_excel(excel_output_path, index=False)

# ================== PARALLEL ==================
def send_to_gemini_for_json(df_chunk: pd.DataFrame, prompt: str, header: str) -> dict:
    print(f'Begin process {df_chunk}')
    table_text = df_chunk.to_csv(index=False)
    full_prompt = f"{prompt}\n\n Below is header and note {header}\nBelow is the table text (CSV):\n{table_text}\nReturn the JSON."
    result_text, _ = run_process_internal_base_v2(
        file_bytes=None,
        filename=None,
        mime=None,
        question=full_prompt,
        model_choice="Gemini 2.5 Flash",
        temperature=0.4,
        top_p=1.0
    )
    #print(f'End process {df_chunk}')
    return safe_parse_json(result_text)

def process_all_chunks_with_threadpool(chunks: dict[str, pd.DataFrame], prompt: str, header: str, max_workers: int = 5) -> list[dict]:
    all_results = []
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(send_to_gemini_for_json, chunk, prompt, header): airline
            for airline, chunk in chunks.items() #if re.match(r"^\\d+", airline.strip())
        }
        for future in as_completed(futures):
            airline = futures[future]
            try:
                result = future.result()
                if result:
                    all_results.extend(result if isinstance(result, list) else [result])
            except Exception as e:
                print(f"❌ Error with {airline}: {e}")
    return all_results

# ================== GEMINI BASE ==================
def run_process_internal_base_v2(file_bytes, filename, mime, question, model_choice, temperature, top_p, batch_size=3):
    api_key = get_next_key()
    genai.configure(api_key=api_key)
    model_name = INTERNAL_MODEL_MAP.get(model_choice, "gemini-2.5-flash")
    print(f'Use key: {api_key}')
    user_prompt = (question or "").strip() or PROMPT_FREIGHT_JSON
    model = genai.GenerativeModel(model_name=model_name,
                                  generation_config={"temperature": float(temperature), "top_p": float(top_p)})
    if file_bytes is None:
        response = model.generate_content(user_prompt)
        #print(response.text)
        return response.text, None
    pages = pdf_to_images(file_bytes)
    all_text_results = []
    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")
                uploaded.append(genai.get_file(up.name))
        resp = model.generate_content([user_prompt] + uploaded)
        all_text_results.append(resp.text if hasattr(resp, "text") else "")
        for up in uploaded:
            try: genai.delete_file(up.name)
            except: pass
    return "\n\n".join(all_text_results), None

# ================== MAIN ROUTER ==================
def run_process(file, question, model_choice, temperature, top_p, external_api_url):
    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)
            all_dfs = []
            saved_header = None
            if check_result > 1:
                print("➡️ PDF có nhiều cột/nhiều trang → dùng pdfplumber extract trước rồi Gemini.")
                
                base_name = os.path.splitext(filename)[0]
                tmp_dir = tempfile.gettempdir()
    
                # 🔁 Ghi file PDF tạm để xử lý
                tmp_pdf_path = os.path.join(tmp_dir, f"{base_name}.pdf")
                with open(tmp_pdf_path, "wb") as f:
                    f.write(file_bytes)
    
                # 🔁 Tạo đường dẫn file Excel
                excel_path = os.path.join(tmp_dir, f"{base_name}.xlsx")
    
                # 🛠 Gọi hàm xử lý
                export_pdf_to_excel(tmp_pdf_path, excel_path)
    
                chunks = split_excel_by_airline_header(excel_path)
                header, _ = run_process_internal_base_v2(
                    file_bytes=file_bytes,
                    filename=filename,
                    mime=mime,
                    question=PROMPT_FREIGHT_HEADER_JSON,
                    model_choice=model_choice,
                    temperature=temperature,
                    top_p=top_p
                )
                print(header)
                chunk_files = [] 
                for airline, df_chunk in chunks.items(): 
                    safe_name = re.sub(r"[^\w\s]", "", airline).replace(" ", "_")
                    print (f'airline : {airline}')
                result = process_all_chunks_with_threadpool(chunks, PROMPT_FREIGHT_JSON, header, 5)
                return json.dumps(result, ensure_ascii=False, indent=2), None
            else: 
                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.")
        

        result_text, _ = 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
        )
        return result_text, None             
    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()