"""Export extracted data to structured formats: JSON, CSV, Excel."""
from __future__ import annotations

import csv
import io
import json

from app.schemas.documents import DocumentDetail


def to_json(detail: DocumentDetail) -> bytes:
    payload = {
        "document": {
            "id": detail.id,
            "filename": detail.filename,
            "doc_type": detail.classification.doc_type if detail.classification else None,
            "num_pages": detail.num_pages,
        },
        "classification": detail.classification.model_dump() if detail.classification else None,
        "summary": detail.summary,
        "extraction": detail.extraction.model_dump() if detail.extraction else None,
        "anomalies": [a.model_dump() for a in detail.anomalies],
    }
    return json.dumps(payload, indent=2, default=str).encode("utf-8")


def _record_rows(extraction) -> tuple[list[str], list[list]]:
    """Flatten records into (header, rows) using the canonical key order."""
    keys = list(extraction.record_keys)
    if not keys:  # derive from the records themselves, preserving first-seen order
        seen = set()
        for rec in extraction.records:
            for f in rec.fields:
                if f.name not in seen:
                    keys.append(f.name)
                    seen.add(f.name)
    rows = []
    for rec in extraction.records:
        m = {f.name: f.value for f in rec.fields}
        rows.append([m.get(k) for k in keys])
    return keys, rows


def to_csv(detail: DocumentDetail) -> bytes:
    buf = io.StringIO()
    w = csv.writer(buf)
    ex = detail.extraction
    # Multi-record docs export as a proper record table; otherwise field/value.
    if ex and ex.records:
        keys, rows = _record_rows(ex)
        w.writerow(keys)
        for row in rows:
            w.writerow(["" if c is None else c for c in row])
    else:
        w.writerow(["field", "value", "confidence"])
        if ex:
            for f in ex.fields:
                w.writerow([f.name, f.value, f"{f.confidence:.2f}"])
    return buf.getvalue().encode("utf-8")


def to_xlsx(detail: DocumentDetail) -> bytes:
    from openpyxl import Workbook

    wb = Workbook()
    ws = wb.active
    ws.title = "Fields"
    ws.append(["Field", "Value", "Confidence"])
    ex = detail.extraction
    if ex:
        for f in ex.fields:
            ws.append([f.name, str(f.value) if f.value is not None else "", round(f.confidence, 2)])

        # Repeated records get their own sheet, one row per record.
        if ex.records:
            keys, rows = _record_rows(ex)
            rec = wb.create_sheet(title=(ex.record_type or "Records")[:31])
            rec.append(keys)
            for row in rows:
                rec.append([str(c) if c is not None else "" for c in row])

        for ti, table in enumerate(ex.tables):
            tab = wb.create_sheet(title=(table.title or f"Table {ti+1}")[:31])
            if table.columns:
                tab.append(table.columns)
            for row in table.rows:
                tab.append([str(c) if c is not None else "" for c in row])

    out = io.BytesIO()
    wb.save(out)
    return out.getvalue()


EXPORTERS = {
    "json": (to_json, "application/json", "json"),
    "csv": (to_csv, "text/csv", "csv"),
    "xlsx": (to_xlsx,
             "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "xlsx"),
}