#!/usr/bin/env python3

"""
Bangladesh BOQ vs SOR Checker (Hugging Face Spaces, Gradio)

- Upload:
  - BOQ PDF
  - Notice PDF (optional)
  - TDS PDF (optional)
  - One SOR PDF (LGED / PWD / BWDB)
- Extract BOQ items from BOQ PDF (simple text-based parser)
- Extract raw text from Notice/TDS PDFs
- Parse SOR PDF (heuristic)
- PWD-first reference logic (strip PWD from code or description, keep numeric)
- District → Zone mapping for LGED/PWD SOR
- Compare BOQ quoted rates with SOR rates
- Show comparison table
- Export: CSV, XLSX, PDF review, DOCX review
"""

import re
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple

import gradio as gr
import pdfplumber
import pandas as pd
import numpy as np

from openpyxl import Workbook
from docx import Document
from reportlab.lib.pagesizes import A4
from reportlab.lib import colors
from reportlab.platypus import (
    SimpleDocTemplate, Table, TableStyle,
    Paragraph, Spacer, PageBreak
)
from reportlab.lib.styles import getSampleStyleSheet

# --------------------------------------------------------------------------------------
# Paths (HF Spaces friendly)
# --------------------------------------------------------------------------------------

BASE_DIR = Path(__file__).parent
UPLOAD_DIR = BASE_DIR / "uploads"
OUTPUT_DIR = BASE_DIR / "outputs"

for d in [UPLOAD_DIR, OUTPUT_DIR]:
    d.mkdir(parents=True, exist_ok=True)

# --------------------------------------------------------------------------------------
# BD-specific Zone Mapping and PWD/LGED Reference Logic
# --------------------------------------------------------------------------------------

BD_ZONE_GROUPS = {
    "A": ["Dhaka", "Mymensingh"],
    "B": ["Chattogram", "Sylhet"],
    "C": ["Rajshahi", "Rangpur"],
    "D": ["Khulna", "Barishal", "Gopalgonj"],
}

NUM_PATTERN = re.compile(r'(\d+(?:[.\-]\d+)*)')
PWD_CODE_PATTERN_DESC = re.compile(r'\bPWD\s*([0-9]+(?:\.[0-9]+)*)\b', re.IGNORECASE)
LGED_CODE_PATTERN_DESC = re.compile(r'\b([0-9]+(?:\.[0-9]+)*)\s*LGED\b', re.IGNORECASE)


def district_to_zone(district: str) -> Optional[str]:
    if not district:
        return None
    d = district.strip().title()
    for zone, districts in BD_ZONE_GROUPS.items():
        if d in districts:
            return zone
    return None


def extract_ref_from_item_code(raw_code: str) -> Tuple[Optional[str], Optional[str]]:
    """
    Priority:
      1) If 'PWD' appears anywhere, treat as PWD SOR item, take last numeric pattern.
      2) Else if 'LGED' appears, treat as LGED SOR item, take last numeric pattern.
      3) Else, (None, None).
    Examples:
      '01.1PWD01.1.3PWD' -> ('PWD', '01.1.3')
      'PWD 03.4.2'       -> ('PWD', '03.4.2')
      '3.12.06LGED'      -> ('LGED', '3.12.06')
    """
    if not raw_code:
        return None, None
    s = str(raw_code).upper()

    if "PWD" in s:
        nums = NUM_PATTERN.findall(s)
        if nums:
            return "PWD", nums[-1]
        return "PWD", None

    if "LGED" in s:
        nums = NUM_PATTERN.findall(s)
        if nums:
            return "LGED", nums[-1]
        return "LGED", None

    return None, None


def extract_ref_from_description(desc: str) -> Tuple[Optional[str], Optional[str]]:
    if not desc:
        return None, None
    m_pwd = PWD_CODE_PATTERN_DESC.search(desc)
    if m_pwd:
        return "PWD", m_pwd.group(1)
    m_lged = LGED_CODE_PATTERN_DESC.search(desc)
    if m_lged:
        return "LGED", m_lged.group(1)
    return None, None


def enrich_boq_with_refs(df: pd.DataFrame) -> pd.DataFrame:
    """
    Add ref_agency, ref_code columns to BOQ DataFrame based on item_code/description.
    """
    agencies = []
    codes = []
    for _, row in df.iterrows():
        code = str(row.get("item_code") or row.get("code") or "")
        desc = str(row.get("description") or row.get("desc") or "")
        agency, ref_code = extract_ref_from_item_code(code)
        if not agency:
            agency, ref_code = extract_ref_from_description(desc)
        agencies.append(agency)
        codes.append(ref_code)
    df = df.copy()
    df["ref_agency"] = agencies
    df["ref_code"] = codes
    return df

# --------------------------------------------------------------------------------------
# Helpers
# --------------------------------------------------------------------------------------

def norm(v) -> str:
    return "" if v is None else str(v).strip()


def to_num(v) -> Optional[float]:
    try:
        return float(str(v).replace(",", "").strip())
    except Exception:
        return None


def read_pdf_text(path: Optional[Path]) -> str:
    if not path:
        return ""
    try:
        text_parts = []
        with pdfplumber.open(str(path)) as pdf:
            for page in pdf.pages:
                text_parts.append(page.extract_text() or "")
        return "\n".join(text_parts)
    except Exception:
        return ""

# --------------------------------------------------------------------------------------
# BOQ Parsing from PDF
# --------------------------------------------------------------------------------------

BOQ_PATTERN = re.compile(
    r"^(\d+)\s+(\S+)\s+(.*?)\s+(\d[\d,]*\.?\d*)\s+([A-Za-z]+)\s+([\d,]*\.?\d+)\s+([\d,]*\.?\d+)$"
)


def parse_boq_pdf(path: Path) -> pd.DataFrame:
    """
    Very simple BOQ line parser from PDF text, expecting lines like:
        item_no code description qty unit rate total
    """
    text = read_pdf_text(path)
    rows: List[Dict[str, Any]] = []

    for line in text.splitlines():
        s = " ".join(line.split())
        m = BOQ_PATTERN.search(s)
        if m:
            item_no, code, desc, qty, unit, rate, total = m.groups()
            rows.append(
                {
                    "item_no": item_no,
                    "item_code": code,
                    "description": desc[:200],
                    "quantity": to_num(qty),
                    "unit": unit,
                    "quoted_rate": to_num(rate),
                    "quoted_amount": to_num(total),
                }
            )

    if not rows:
        # Fallback example when parsing fails
        rows = [
            {
                "item_no": "1",
                "item_code": "04-120",
                "description": "Construction of B.M. Pillars",
                "quantity": 2.0,
                "unit": "Nos",
                "quoted_rate": 1412.58,
                "quoted_amount": 2825.16,
            }
        ]

    return pd.DataFrame(rows)

# --------------------------------------------------------------------------------------
# SOR Parsing from PDF (BWDB-style)
# --------------------------------------------------------------------------------------

def parse_sor_pdf(path: Path, default_agency: str = "BWDB") -> pd.DataFrame:
    """
    Heuristic SOR parser from PDF text for BWDB-style codes (xx-xxx[-xx]).
    Expected rough line shape:
        code description ... unit rate
    """
    text = read_pdf_text(path)
    rows: List[Dict[str, Any]] = []

    for line in text.splitlines():
        s = " ".join(line.split())
        if not s:
            continue

        m = re.match(r"^(\d{2}-\d{3}(?:-\d{2})?)\s+(.*)$", s)
        if m:
            code = m.group(1)
            rest = m.group(2)
            tokens = rest.split()
            rate_token = None
            rate_idx = None
            for i in range(len(tokens) - 1, -1, -1):
                if to_num(tokens[i]) is not None:
                    rate_token = tokens[i]
                    rate_idx = i
                    break
            if rate_token is None or rate_idx is None or rate_idx == 0:
                continue
            unit = tokens[rate_idx - 1]
            desc_tokens = tokens[: rate_idx - 1]
            desc = " ".join(desc_tokens)
            rows.append(
                {
                    "agency": default_agency,
                    "code": code,
                    "ref_code": None,
                    "description": desc,
                    "unit": unit,
                    "rate": to_num(rate_token),
                    "zone": None,
                }
            )

    return pd.DataFrame(rows)

# --------------------------------------------------------------------------------------
# SOR Parsing from PDF (LGED/PWD-style)
# --------------------------------------------------------------------------------------

def parse_lged_pwd_sor_pdf(path: Path, agency_hint: Optional[str] = None) -> pd.DataFrame:
    """
    Heuristic parser for LGED/PWD SOR PDF:
    - Tries to detect item codes (01.1.3, 3.12.06, etc.)
    - Uses last numeric as rate
    - Unit is token before rate
    - Uses PWD/LGED pattern in description to set ref_code
    """
    text = read_pdf_text(path)
    rows: List[Dict[str, Any]] = []

    for line in text.splitlines():
        s = " ".join(line.split())
        if not s:
            continue

        m = re.match(r"^(\d+(?:\.\d+)+)\s+(.*)$", s)
        if m:
            code = m.group(1)
            rest = m.group(2)
            tokens = rest.split()
            rate_token = None
            rate_idx = None
            for i in range(len(tokens) - 1, -1, -1):
                if to_num(tokens[i]) is not None:
                    rate_token = tokens[i]
                    rate_idx = i
                    break
            if rate_token is None or rate_idx is None or rate_idx == 0:
                continue

            unit = tokens[rate_idx - 1]
            desc_tokens = tokens[: rate_idx - 1]
            desc = " ".join(desc_tokens)

            ref_agency, ref_code = extract_ref_from_description(desc)
            agency = ref_agency or (agency_hint or "LGED")

            rows.append(
                {
                    "agency": agency,
                    "code": code,
                    "ref_code": ref_code,
                    "description": desc,
                    "unit": unit,
                    "rate": to_num(rate_token),
                    "zone": None,
                }
            )

    return pd.DataFrame(rows)

# --------------------------------------------------------------------------------------
# Matching BOQ vs SOR
# --------------------------------------------------------------------------------------

def match_boq_to_sor(
    boq_df: pd.DataFrame,
    sor_bwdb: pd.DataFrame,
    sor_lged_pwd: pd.DataFrame,
    project_district: str,
) -> pd.DataFrame:
    """
    Matching strategy:
      - Direct BWDB code match (item_code == SOR code)
      - Else PWD/LGED ref match: (ref_agency, ref_code) against LGED/PWD SOR
    Zone is determined from project_district and attached as info.
    """
    zone = district_to_zone(project_district)

    # Index BWDB by code
    bwdb_map = {}
    if not sor_bwdb.empty:
        for _, r in sor_bwdb.iterrows():
            c = str(r["code"]).strip()
            bwdb_map[c] = r

    # Index LGED/PWD by (agency, ref_code)
    lged_pwd_map = {}
    if not sor_lged_pwd.empty:
        for _, r in sor_lged_pwd.iterrows():
            agency = str(r.get("agency") or "").upper()
            ref_code = str(r.get("ref_code") or "").strip()
            if agency and ref_code:
                lged_pwd_map[(agency, ref_code)] = r

    # Enrich BOQ with PWD/LGED references
    boq_df = enrich_boq_with_refs(boq_df)

    sor_rate_col = []
    sor_agency_col = []
    sor_code_col = []
    sor_ref_code_col = []
    sor_zone_col = []
    diff_col = []
    pct_diff_col = []
    flag_col = []

    for _, row in boq_df.iterrows():
        code = str(row.get("item_code") or "").strip()
        agency = str(row.get("ref_agency") or "").upper()
        ref_code = str(row.get("ref_code") or "").strip()
        boq_rate = row.get("quoted_rate")

        sor_rate = None
        sor_agency = None
        sor_code = None
        sor_ref = None
        sor_zone = zone

        # 1) direct BWDB match
        if code and code in bwdb_map:
            r = bwdb_map[code]
            sor_rate = r["rate"]
            sor_agency = r.get("agency", "BWDB")
            sor_code = r.get("code")
            sor_ref = r.get("ref_code")

        # 2) PWD/LGED ref match
        if sor_rate is None and agency and ref_code:
            r = lged_pwd_map.get((agency, ref_code))
            if r is not None:
                sor_rate = r["rate"]
                sor_agency = r.get("agency")
                sor_code = r.get("code")
                sor_ref = r.get("ref_code")

        sor_rate_col.append(sor_rate)
        sor_agency_col.append(sor_agency)
        sor_code_col.append(sor_code)
        sor_ref_code_col.append(sor_ref)
        sor_zone_col.append(sor_zone)

        diff = None
        pct = None
        flag = "OK"

        if sor_rate is None:
            flag = "SOR missing"
        elif boq_rate is None:
            flag = "BOQ rate missing"
        else:
            diff = round(boq_rate - sor_rate, 2)
            if sor_rate:
                pct = round((diff / sor_rate) * 100, 2)
            if pct is not None:
                if abs(pct) > 10:
                    flag = "MISMATCH"
                elif abs(pct) > 0:
                    flag = "VARIANCE"

        diff_col.append(diff)
        pct_diff_col.append(pct)
        flag_col.append(flag)

    out = boq_df.copy()
    out["project_district"] = project_district
    out["project_zone"] = zone
    out["sor_rate"] = sor_rate_col
    out["sor_agency"] = sor_agency_col
    out["sor_code"] = sor_code_col
    out["sor_ref_code"] = sor_ref_code_col
    out["sor_zone"] = sor_zone_col
    out["diff"] = diff_col
    out["pct_diff"] = pct_diff_col
    out["flag"] = flag_col

    return out

# --------------------------------------------------------------------------------------
# Export functions: CSV, XLSX, PDF, DOCX
# --------------------------------------------------------------------------------------

def export_to_csv(df: pd.DataFrame, tender_id: str) -> str:
    out = OUTPUT_DIR / f"{tender_id}_boq_sor_diff.csv"
    df.to_csv(out, index=False)
    return str(out)


def export_to_xlsx(df: pd.DataFrame, tender_id: str) -> str:
    out = OUTPUT_DIR / f"{tender_id}_boq_sor_diff.xlsx"
    wb = Workbook()
    ws = wb.active
    ws.title = "BOQ SOR Diff"
    headers = list(df.columns)
    ws.append(headers)
    for _, row in df.iterrows():
        ws.append([row.get(col) for col in headers])
    wb.save(out)
    return str(out)


def export_to_pdf(df: pd.DataFrame, tender_id: str) -> str:
    out = OUTPUT_DIR / f"{tender_id}_review_sheet.pdf"
    doc = SimpleDocTemplate(
        str(out),
        pagesize=A4,
        rightMargin=24,
        leftMargin=24,
        topMargin=24,
        bottomMargin=24,
    )
    styles = getSampleStyleSheet()
    elems = [
        Paragraph(f"BOQ vs SOR Review Sheet - {tender_id}", styles["Title"]),
        Spacer(1, 8),
        Paragraph(
            "Internal comparison between BOQ quoted rates and SOR rates.",
            styles["Normal"],
        ),
        Spacer(1, 12),
    ]

    data = [
        [
            "Item",
            "Code",
            "Description",
            "Qty",
            "Unit",
            "BOQ Rate",
            "SOR Rate",
            "Agency",
            "Zone",
            "Diff",
            "% Diff",
            "Flag",
        ]
    ]
    for _, r in df.iterrows():
        data.append(
            [
                str(r.get("item_no", "")),
                str(r.get("item_code", "")),
                str(r.get("description", ""))[:60],
                str(r.get("quantity", "")),
                str(r.get("unit", "")),
                str(r.get("quoted_rate", "")),
                str(r.get("sor_rate", "")),
                str(r.get("sor_agency", "")),
                str(r.get("sor_zone", "")),
                str(r.get("diff", "")),
                str(r.get("pct_diff", "")),
                str(r.get("flag", "")),
            ]
        )

    tbl = Table(data, repeatRows=1)
    tbl.setStyle(
        TableStyle(
            [
                ("BACKGROUND", (0, 0), (-1, 0), colors.HexColor("#d9eaf7")),
                ("GRID", (0, 0), (-1, -1), 0.25, colors.grey),
                ("FONTSIZE", (0, 0), (-1, -1), 7),
                ("VALIGN", (0, 0), (-1, -1), "TOP"),
            ]
        )
    )

    elems.append(tbl)
    elems.append(PageBreak())
    elems.append(Paragraph("Notes", styles["Heading2"]))
    elems.append(
        Paragraph("1. Review all MISMATCH and VARIANCE rows before submission.", styles["Normal"])
    )
    elems.append(
        Paragraph("2. Confirm correct zone selection based on project district.", styles["Normal"])
    )
    elems.append(
        Paragraph("3. Ensure all SOR references (PWD/LGED) are correctly interpreted.", styles["Normal"])
    )

    doc.build(elems)
    return str(out)


def export_to_docx(df: pd.DataFrame, tender_id: str) -> str:
    out = OUTPUT_DIR / f"{tender_id}_review_sheet.docx"
    doc = Document()
    doc.add_heading(f"BOQ vs SOR Review Sheet - {tender_id}", level=1)
    doc.add_paragraph("Internal comparison between BOQ quoted rates and SOR rates.")

    headers = [
        "Item",
        "Code",
        "Description",
        "Qty",
        "Unit",
        "BOQ Rate",
        "SOR Rate",
        "Agency",
        "Zone",
        "Diff",
        "% Diff",
        "Flag",
    ]
    table = doc.add_table(rows=1, cols=len(headers))
    hdr_cells = table.rows[0].cells
    for i, h in enumerate(headers):
        hdr_cells[i].text = h

    for _, r in df.iterrows():
        row_cells = table.add_row().cells
        row_cells[0].text = str(r.get("item_no", ""))
        row_cells[1].text = str(r.get("item_code", ""))
        row_cells[2].text = str(r.get("description", ""))[:120]
        row_cells[3].text = str(r.get("quantity", ""))
        row_cells[4].text = str(r.get("unit", ""))
        row_cells[5].text = str(r.get("quoted_rate", ""))
        row_cells[6].text = str(r.get("sor_rate", ""))
        row_cells[7].text = str(r.get("sor_agency", ""))
        row_cells[8].text = str(r.get("sor_zone", ""))
        row_cells[9].text = str(r.get("diff", ""))
        row_cells[10].text = str(r.get("pct_diff", ""))
        row_cells[11].text = str(r.get("flag", ""))

    doc.add_page_break()
    doc.add_heading("Notes", level=2)
    doc.add_paragraph("1. Review all MISMATCH and VARIANCE rows before submission.")
    doc.add_paragraph("2. Confirm correct zone selection based on project district.")
    doc.add_paragraph("3. Ensure all SOR references (PWD/LGED) are correctly interpreted.")

    doc.save(out)
    return str(out)

# --------------------------------------------------------------------------------------
# Main Gradio Pipeline
# --------------------------------------------------------------------------------------

def process_pipeline(
    boq_pdf: Optional[gr.File],
    notice_pdf: Optional[gr.File],
    tds_pdf: Optional[gr.File],
    sor_pdf: Optional[gr.File],
    project_district: str,
    tender_id: str,
):
    if not boq_pdf:
        return "Please upload BOQ PDF.", None, None, None, None, None

    tender_id = tender_id or "TENDER"

    # Convert gr.File to Path
    local_boq = Path(boq_pdf.name)
    local_notice = Path(notice_pdf.name) if notice_pdf else None
    local_tds = Path(tds_pdf.name) if tds_pdf else None
    local_sor = Path(sor_pdf.name) if sor_pdf else None

    # Parse BOQ
    boq_df = parse_boq_pdf(local_boq)

    # Parse Notice/TDS text (currently not used for matching)
    _ = read_pdf_text(local_notice)
    _ = read_pdf_text(local_tds)

    # Parse SOR (single PDF)
    sor_bwdb = pd.DataFrame()
    sor_lged_pwd = pd.DataFrame()
    if local_sor:
        name = local_sor.name.lower()
        if "lged" in name or "pwd" in name:
            sor_lged_pwd = parse_lged_pwd_sor_pdf(local_sor)
        else:
            sor_bwdb = parse_sor_pdf(local_sor, default_agency="BWDB")

    # Compare
    result_df = match_boq_to_sor(boq_df, sor_bwdb, sor_lged_pwd, project_district)

    # Exports
    csv_path = export_to_csv(result_df, tender_id)
    xlsx_path = export_to_xlsx(result_df, tender_id)
    pdf_path = export_to_pdf(result_df, tender_id)
    docx_path = export_to_docx(result_df, tender_id)

    zone = result_df["project_zone"].iloc[0] if "project_zone" in result_df.columns else "N/A"
    status = f"Completed. Items: {len(result_df)}, Zone: {zone}"

    return (
        status,
        result_df,
        csv_path,
        xlsx_path,
        pdf_path,
        docx_path,
    )

# --------------------------------------------------------------------------------------
# Gradio UI
# --------------------------------------------------------------------------------------

with gr.Blocks() as demo:
    gr.Markdown("## Bangladesh BOQ vs SOR Checker (PDF-based)")

    with gr.Row():
        boq_pdf = gr.File(label="BOQ PDF", file_types=[".pdf"])
        notice_pdf = gr.File(label="Notice PDF (optional)", file_types=[".pdf"])
        tds_pdf = gr.File(label="TDS PDF (optional)", file_types=[".pdf"])

    sor_pdf = gr.File(
        label="SOR PDF (LGED/PWD/BWDB)",
        file_types=[".pdf"],
    )

    with gr.Row():
        tender_id_in = gr.Textbox(label="Tender ID", value="ML-PW-41")
        district_in = gr.Textbox(label="Project District (e.g. Barishal)", value="Barishal")

    run_btn = gr.Button("Run Analysis")

    status_out = gr.Textbox(label="Status")
    table_out = gr.Dataframe(label="BOQ vs SOR Result", wrap=True)

    csv_out = gr.File(label="Download CSV")
    xlsx_out = gr.File(label="Download Excel")
    pdf_out = gr.File(label="Download PDF")
    docx_out = gr.File(label="Download Word")

    run_btn.click(
        fn=process_pipeline,
        inputs=[boq_pdf, notice_pdf, tds_pdf, sor_pdf, district_in, tender_id_in],
        outputs=[status_out, table_out, csv_out, xlsx_out, pdf_out, docx_out],
    )

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