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 metadata, 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."""
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}')
    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(question)
        #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([question] + 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)
        check_result = check_pdf_structure(file_bytes)
        if check_result > 1:
            base_name = os.path.splitext(filename)[0]
            tmp_dir = tempfile.gettempdir()
            excel_path = os.path.join(tmp_dir, f"{base_name}.xlsx")
            export_pdf_to_excel(filename, 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:
            return "Only supports multi-airline PDF for now", 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()