"""
PDF to Excel table extraction API.
FastAPI version - No Gradio needed.
"""

from fastapi import FastAPI, File, UploadFile, Form
from fastapi.middleware.cors import CORSMiddleware
from markitdown import MarkItDown as DocumentConverter
import pandas as pd
import pdfplumber
import tempfile
import os
import base64
from ocr_support import OcrError, prepare_for_extraction, normalise_cjk, validate_ocr_language
import PyPDF2
import re

app = FastAPI(title="PDF to Excel Converter API")

# Enable CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

md_converter = DocumentConverter()
MIN_TEXT_CHARS_PER_PAGE = 40

def clean_text(text):
    """Remove NULL bytes and control characters"""
    if not text:
        return text
    # Remove NULL bytes
    text = text.replace('\x00', '')
    # Remove other control characters except newline, tab, and carriage return
    text = re.sub(r'[\x01-\x08\x0B\x0C\x0E-\x1F\x7F-\x9F]', '', text)
    return text

def normalize_cell(value):
    """Normalize table cell values from PDF extraction."""
    if value is None:
        return ""
    return normalise_cjk(clean_text(str(value))).strip()

def dedupe_headers(headers):
    """Ensure spreadsheet column names are present and unique."""
    used = {}
    output = []

    for index, header in enumerate(headers):
        base = normalize_cell(header) or f"Column_{index + 1}"
        count = used.get(base, 0)
        used[base] = count + 1
        output.append(base if count == 0 else f"{base}_{count + 1}")

    return output

def dataframe_from_rows(rows):
    """Convert extracted raw rows into a clean DataFrame."""
    cleaned_rows = [
        [normalize_cell(cell) for cell in row]
        for row in rows
        if row and any(normalize_cell(cell) for cell in row)
    ]

    if len(cleaned_rows) < 2:
        return None

    max_width = max(len(row) for row in cleaned_rows)
    normalized_rows = [row + [""] * (max_width - len(row)) for row in cleaned_rows]

    # Drop columns that are empty from top to bottom.
    useful_columns = [
        index
        for index in range(max_width)
        if any(row[index] for row in normalized_rows)
    ]
    normalized_rows = [[row[index] for index in useful_columns] for row in normalized_rows]

    if not normalized_rows or len(normalized_rows[0]) < 2:
        return None

    headers = dedupe_headers(normalized_rows[0])
    data_rows = normalized_rows[1:]

    if not any(any(cell for cell in row) for row in data_rows):
        return None

    return pd.DataFrame(data_rows, columns=headers)


def extract_tables_with_pdfplumber(pdf_path):
    tables, text_chars = [], 0
    reader = PyPDF2.PdfReader(pdf_path)
    with pdfplumber.open(pdf_path) as pdf, tempfile.TemporaryDirectory() as workdir:
        for index, page in enumerate(pdf.pages):
            text_chars += len(page.extract_text() or '')
            page_tables = []
            for candidate in page.find_tables(table_settings={
                'vertical_strategy': 'lines', 'horizontal_strategy': 'lines',
                'snap_tolerance': 3, 'join_tolerance': 3, 'intersection_tolerance': 3,
            }):
                frame = dataframe_from_rows(candidate.extract())
                if frame is not None:
                    page_tables.append(frame)
            if not page_tables:
                page_path = os.path.join(workdir, f'page-{index + 1}.pdf')
                writer = PyPDF2.PdfWriter()
                writer.add_page(reader.pages[index])
                with open(page_path, 'wb') as stream:
                    writer.write(stream)
                result = md_converter.convert(page_path)
                page_tables = extract_tables_from_markdown(clean_text(result.text_content or ''))
            tables.extend(page_tables)
    return tables, text_chars

def split_markdown_row(line):
    """Keep escaped pipes inside cells, and preserve intentionally empty cells."""
    text = line.strip()
    cells, current = [], []
    index = 0
    while index < len(text):
        if text[index:index + 2] == r'\|':
            current.append('|')
            index += 2
        elif text[index] == '|':
            cells.append(''.join(current).strip())
            current = []
            index += 1
        else:
            current.append(text[index])
            index += 1
    cells.append(''.join(current).strip())
    if text.startswith('|'):
        cells = cells[1:]
    if text.endswith('|') and cells and not cells[-1]:
        cells = cells[:-1]
    return cells

def extract_tables_from_markdown(markdown_text):
    """Require a real Markdown header separator before treating text as a table."""
    lines = markdown_text.splitlines()
    tables = []
    index = 0
    while index + 1 < len(lines):
        separator = split_markdown_row(lines[index + 1])
        is_separator = separator and all(re.fullmatch(r':?-{3,}:?', cell) for cell in separator)
        if '|' not in lines[index] or not is_separator:
            index += 1
            continue
        end = index + 2
        while end < len(lines) and lines[end].strip() and '|' in lines[end]:
            end += 1
        table = parse_markdown_table(lines[index:end])
        if table is not None and not table.empty:
            tables.append(table)
        index = end
    return tables

def parse_markdown_table(table_lines):
    if len(table_lines) < 3:
        return None
    rows = [split_markdown_row(line) for line in table_lines]
    separator = rows[1]
    if not separator or not all(re.fullmatch(r':?-{3,}:?', cell) for cell in separator):
        return None
    # The padding helper keeps extra cells rather than silently dropping a row.
    return dataframe_from_rows([rows[0], *rows[2:]])

def estimate_text_chars(pdf_path):
    """Estimate whether a PDF contains selectable text."""
    text_chars = 0
    with open(pdf_path, 'rb') as f:
        pdf_reader = PyPDF2.PdfReader(f)
        for page in pdf_reader.pages:
            text_chars += len(page.extract_text() or "")
    return text_chars

def format_spreadsheet(sheet, frame):
    """Make amount/quantity columns usable while keeping IDs and dates as text."""
    from openpyxl.styles import Alignment
    from openpyxl.utils import get_column_letter

    for column_index, header in enumerate(frame.columns, start=1):
        label = normalise_cjk(str(header))
        identifier = bool(re.search(r'\b(?:id|code|reference|account|date|phone|postcode)\b|编号|編號|代码|代碼|税号|稅號|日期|账号|帳號|郵編|邮编|电话|電話', label, re.I))
        numeric_column = not identifier and bool(re.search(r'quantity|\bqty\b|price|amount|total|balance|credit|debit|\bgst\b|\btax\b|cost|rate|hours|数量|數量|单价|單價|金额|金額|总计|總計|合计|合計|小计|小計|税额|稅額|余额|餘額|成本|费率|費率|工时|工時', label, re.I))
        for row in range(2, sheet.max_row + 1):
            cell = sheet.cell(row, column_index)
            value = cell.value
            if isinstance(value, str) and value.startswith('='):
                cell.data_type = 's'
            if not numeric_column or not isinstance(value, str):
                continue
            text = value.strip()
            bracketed = text.startswith('(') and text.endswith(')')
            if bracketed:
                text = text[1:-1].strip()
            match = re.fullmatch(r'(?P<currency>[A-Z]{3}\s*|[$€£¥￥]\s*)?(?P<number>[+-]?(?:\d{1,3}(?:,\d{3})+|\d+)(?:\.\d+)?)(?P<percent>%)?', text)
            if not match or re.fullmatch(r'[+-]?0\d+', match['number']):
                continue
            if len(re.sub(r'\D', '', match['number']).lstrip('0')) > 15:
                continue
            number = float(match['number'].replace(',', ''))
            if bracketed:
                number = -abs(number)
            decimals = len(match['number'].split('.')[-1]) if '.' in match['number'] else 0
            decimals = min(decimals, 8)
            if match['percent']:
                cell.value = number / 100
                cell.number_format = '0' + ('.' + '0' * decimals if decimals else '') + '%'
            else:
                cell.value = number
                prefix = '"' + match['currency'].strip() + '" ' if match['currency'] else ''
                cell.number_format = prefix + '#,##0' + ('.' + '0' * decimals if decimals else '')
                if bracketed:
                    cell.number_format += ';(' + cell.number_format + ')'
            cell.alignment = Alignment(horizontal='right', vertical='top')
        width = max(len(str(header)), *(min(60, len(str(value))) for value in frame.iloc[:, column_index - 1]))
        sheet.column_dimensions[get_column_letter(column_index)].width = min(60, max(12, width + 2))
    sheet.freeze_panes = 'A2'

@app.post("/call/convert")
async def convert_pdf_to_excel(file: UploadFile = File(...), output_format: str = Form("xlsx"), enable_ocr: bool = Form(False), ocr_language: str = Form("eng")):
    """Convert PDF to Excel - compatible with Gradio API format"""

    pdf_path = None
    output_path = None
    try:
        language = validate_ocr_language(ocr_language if isinstance(ocr_language, str) else 'eng')
        # Save uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_pdf:
            content = await file.read()
            tmp_pdf.write(content)
            pdf_path = tmp_pdf.name

        # Check page limit
        with open(pdf_path, 'rb') as f:
            pdf_reader = PyPDF2.PdfReader(f)
            num_pages = len(pdf_reader.pages)

        if num_pages > 10:
            os.remove(pdf_path)
            return {
                "event_id": "error",
                "message": f"📄 Your PDF has {num_pages} pages. Our free plan converts up to 10 pages at a time.\n\n💎 Want unlimited pages? Pro plan coming soon! Split your PDF or try a smaller file for now."
            }

        # Try direct table extraction first. This is stronger for many
        # text-based invoices, bank statements, and reports than Markdown.
        with prepare_for_extraction(pdf_path, require_table=True, enabled=enable_ocr is True, ocr_language=language) as (prepared, ocr_count):
            tables, text_chars = extract_tables_with_pdfplumber(prepared)

        if text_chars < max(MIN_TEXT_CHARS_PER_PAGE * num_pages, 1):
            os.remove(pdf_path)
            return {
                "event_id": "error",
                "message": "Not enough readable text was found. Use a clearer, upright scan in the selected OCR language or a PDF with selectable text."
            }

        if not tables:
            os.remove(pdf_path)
            return {
                "event_id": "error",
                "message": "No convertible tables were found.\n\nThis tool works best with text-based PDFs that contain clear rows and columns. Borderless tables, merged cells, and heavily designed reports may need manual cleanup."
            }

        # Convert to Excel or CSV
        if output_format == "xlsx":
            output_path = tempfile.NamedTemporaryFile(delete=False, suffix='.xlsx').name
            with pd.ExcelWriter(output_path, engine='openpyxl') as writer:
                for i, df in enumerate(tables):
                    sheet_name = f'Table_{i+1}'
                    df.to_excel(writer, sheet_name=sheet_name, index=False)
                    format_spreadsheet(writer.sheets[sheet_name], df)
        else:  # csv
            output_path = tempfile.NamedTemporaryFile(delete=False, suffix='.csv').name
            if len(tables) == 1:
                tables[0].to_csv(output_path, index=False)
            else:
                if any(list(table.columns) != list(tables[0].columns) for table in tables[1:]):
                    raise ValueError('The detected tables have different columns. Choose XLSX to keep each table separately, then export the sheet you need as CSV.')
                combined_df = pd.concat(tables, ignore_index=True)
                combined_df.to_csv(output_path, index=False)

        # Read file and convert to base64
        with open(output_path, 'rb') as f:
            file_bytes = f.read()
            base64_data = base64.b64encode(file_bytes).decode('utf-8')

        # Cleanup
        os.remove(pdf_path)
        os.remove(output_path)

        # Return in Gradio-compatible format
        extension = '.xlsx' if output_format == 'xlsx' else '.csv'
        return {
            "event_id": "success",
            "file_data": base64_data,
            "filename": file.filename.replace('.pdf', extension),
            "table_count": len(tables),
            "row_count": int(sum(len(df.index) for df in tables)),
            "ocr_page_count": ocr_count,
            "ocr_language": language
        }

    except Exception as e:
        error_msg = str(e)

        # Provide user-friendly error messages
        if isinstance(e, OcrError):
            friendly_msg = str(e)
        elif 'needs OCR' in error_msg or 'different columns' in error_msg:
            friendly_msg = error_msg
        elif "XML compatible" in error_msg or "NULL bytes" in error_msg or "encoding" in error_msg.lower():
            friendly_msg = "😔 Sorry! Your PDF contains special characters that we couldn't process.\n\n✨ Try:\n• Re-saving your PDF from the original document\n• Using a different PDF viewer to export it\n• Contact us if the issue persists"
        elif "password" in error_msg.lower() or "encrypted" in error_msg.lower():
            friendly_msg = "🔒 This PDF is password-protected.\n\n💡 Please remove the password first, then try again!"
        elif "corrupt" in error_msg.lower() or "damaged" in error_msg.lower():
            friendly_msg = "⚠️ Your PDF file appears to be corrupted.\n\n🔧 Try:\n• Re-downloading the original PDF\n• Opening and re-saving it in Adobe Reader"
        else:
            friendly_msg = f"😔 Oops! Something went wrong during conversion.\n\n🛠️ Technical details: {error_msg}\n\n💡 Try uploading a different PDF or contact support if this persists."

        return {
            "event_id": "error",
            "message": friendly_msg
        }
    finally:
        for path in (pdf_path, output_path):
            if path and os.path.isfile(path):
                os.remove(path)

@app.get("/")
async def root():
    return {
        "message": "PDF to Excel Converter API",
        "powered_by": "table extraction pipeline",
        "endpoint": "/call/convert"
    }

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