"""Run the figure workflow end to end.

The workflow is independent of the rest of the pipeline. It takes a parse and
the source document, and gives back an enriched copy of that parse, plus the
page scans. Nothing about it depends on being called from inside a larger run,
so a parse produced earlier can be enriched without reproducing it, and
sequencing is a matter of calling the steps in order rather than a flag inside
a monolithic run.

The enriched parse is written under the filename it was read from, so pointing
`output_dir` at the parse's own directory replaces it in place. The parse is an
input every later stage reads by name, and enriching it under that name means
no stage has to be told which of two files to pick up. Run this after fusion,
which is what makes the chart values reach the document the rest of the
pipeline reads.

Pages are nominated from the parse and the PDF's text layer (`nominate`): a
page holding a picture the parse did not class as page furniture, or a table
with no usable text layer under it. Every nominated page is scanned whole, and
no judgement sits between nomination and scan. `pages="all"` adds every other
page of the document, for documents whose content the parse cannot be trusted
to point at. The scans then feed a copy of the parse in the order `run_figures`
sets and explains: tables, figures, and the removal of text a figure was read
over. Figure values are read last, when the caller supplies the page cells and
both value readers, and go to their own artifact, not into the parse.

A page is scanned once when `output_dir` is local. Its stored response is what
a later run reads instead of paying to scan the page again, provided the
stored response's format matches the requested model's generation. A
mismatched one is moved aside rather than deleted. With an `s3://`
`output_dir` no stored response is found, so every page is scanned and billed
again and the upload overwrites the stored response.

The model selects the reading path. A dpt-3 model's responses are digested,
grafted and footnoted by the `dpt3` package. The dpt-3 path also reads any
other page whose dpt-3 response is already stored, and inserts grids the scan
found over empty regions as new tables. Any other model runs the original
path.

Every nominated page reaches the output, as a page scan or, when its scan
failed, as an error. A chart is often where a filing states a number that
appears nowhere else, so a page that was nominated and produced nothing is
exactly what someone needs to see.
"""

from __future__ import annotations

import asyncio
import json
import re
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from docling_core.types.doc.document import DoclingDocument
from loguru import logger

from quber.agents.cell_reader import CellReader, get_cell_reader
from quber.agents.figure_correction import FigureCorrector, get_figure_corrector
from quber.agents.llm_client import LLMClient, get_llm_client
from quber.agents.status_inspector import StatusInspector, get_status_inspector
from quber.core.extractors.base import ExtractedTable
from quber.core.extractors.set_of_mark.assemble import TableAssembly
from quber.core.figures import dpt3
from quber.core.figures.capture import CAPTURE_DPI, capture_tables
from quber.core.figures.correct import correct_figures
from quber.core.figures.graft import graft_figures, graft_tables, unread_pictures
from quber.core.figures.models import FigureRun, FigureValueRun, PageScan
from quber.core.figures.nominate import NominatedPage, nominate_pages, table_verdicts
from quber.core.figures.scan import SCAN_DEFAULT_MODEL, page_scan, scan_page
from quber.core.figures.values import LocalValueReader, ScanValueParser, read_figure_values
from quber.files.output import output_sink, write_completion_marker
from quber.files.pdf import render_page
from quber.providers.landing.client import is_v2_model

#: Pages scanned at once. The scans are independent, and each spends nearly all
#: its time waiting on the platform.
SCAN_CONCURRENCY = 6

#: Table corrections in flight at once, matching the table pipeline's own bound.
CORRECT_CONCURRENCY = 16

#: The two page breadths a run can take. `nominated` scans the pages the parse
#: points at; `all` scans every page of the document, because a page whose
#: content the parse never classified carries no nomination signal at all.
PAGE_SCOPES = ("nominated", "all")


@dataclass
class ChartOutput:
    """Where one figure run's artifacts landed, and what the run produced."""

    run: FigureRun
    document: DoclingDocument
    tables: List[ExtractedTable] = field(default_factory=list)
    artifacts: List[str] = field(default_factory=list)
    values: Optional[FigureValueRun] = None
    digests: Dict[int, dpt3.PageDigest] = field(default_factory=dict)


async def run_figures(
    source: Path,
    document: DoclingDocument,
    base: str,
    parse_name: str,
    output_dir: str = "output",
    model: str = SCAN_DEFAULT_MODEL,
    concurrency: int = SCAN_CONCURRENCY,
    pages: str = "nominated",
    llm: Optional[LLMClient] = None,
    inspector: Optional[StatusInspector] = None,
    cell_reader: Optional[CellReader] = None,
    corrector: Optional[FigureCorrector] = None,
    cells_pages: Optional[Dict[int, dict]] = None,
    value_parser: Optional[ScanValueParser] = None,
    value_reader: Optional[LocalValueReader] = None,
) -> ChartOutput:
    """Read the page scans into a copy of `document`; write the scans and the refined parse.

    `document` itself is not changed. The refined copy is returned as
    `ChartOutput.document` and written under `parse_name`, the filename the
    parse was read from. The parse is an input the rest of the pipeline reads
    by name, and this step enriches it rather than producing a variant, so the
    name never changes and no later step has to be told which file to pick up.
    Pointing `output_dir` at the directory the parse came from therefore
    replaces it. Pointing it elsewhere leaves the original alone.

    Each page's raw response is written as soon as it is scanned. A page whose
    scan fails is reported as an error and the run goes on, so the refined
    parse is then written from the pages that did scan. A step that raises
    after the scans stops the run before the run record or the parse is
    written.
    """
    if pages not in PAGE_SCOPES:
        raise ValueError(f"pages must be one of {PAGE_SCOPES}, not {pages!r}")
    nominated = nominate_pages(document, table_verdicts(document, source))
    logger.info(
        "Figure run: {} nominated {} page(s) of {}: {}",
        base,
        len(nominated),
        len(document.pages),
        [n.page for n in nominated],
    )
    if pages == "all":
        nominated = widen_to_all_pages(document, nominated)
    track = is_v2_model(model)
    if track:
        # A page nothing nominated is read too when its dpt-3 response is
        # already stored. A stat-panel page holds no picture and is never
        # nominated, but its stored response costs nothing to digest and its
        # block partition is what groups the page's text.
        stored = _stored_track_pages(base, output_dir, {n.page for n in nominated})
        if stored:
            logger.info(
                "Figure run: {} reading {} page(s) with a stored dpt-3 response: {}",
                base,
                len(stored),
                [n.page for n in stored],
            )
            nominated = sorted([*nominated, *stored], key=lambda n: n.page)
    llm = llm if llm is not None else get_llm_client(None)
    inspector = inspector if inspector is not None else get_status_inspector()
    cell_reader = cell_reader if cell_reader is not None else get_cell_reader()
    corrector = corrector if corrector is not None else get_figure_corrector()

    scans, digests, errors = await _scan_pages(source, base, nominated, model, output_dir, concurrency)
    scans.sort(key=lambda s: s.page)

    run = FigureRun(document=base, scans=scans, errors=errors)
    dims = _page_dims(document)
    # Tables first: a captured grid replaces the body of a table the parse
    # holds or is added beside a picture, while the figures reshape the
    # pictures around it.
    refined = document.model_copy(deep=True)
    assembly = TableAssembly(source, llm, asyncio.Semaphore(CORRECT_CONCURRENCY), CAPTURE_DPI, inspector)
    tables, capture_errors = await capture_tables(
        refined, scans, source, dims, assembly, cell_reader, orphans=track
    )
    run.errors.extend(capture_errors)
    run.errors.extend(graft_tables(refined, scans, tables, dims))
    if track:
        # A captured grid over a region the document holds nothing for — a
        # map's legend — is inserted at its reading-order position.
        run.errors.extend(dpt3.insert_orphan_tables(refined, scans, tables, dims))

    if track:
        refined, graft_errors = dpt3.graft_figures(refined, scans, dims, source=source)
    else:
        refined, graft_errors = graft_figures(refined, scans, dims, source=source)
    run.errors.extend(graft_errors)
    # After the graft, so it sees only what the picture's own children did not
    # already account for.
    run.removed = await correct_figures(refined, scans, dims, source, corrector)
    if track:
        # After the correction, which records each figure's markers and notes on
        # its picture — the resolution's inputs.
        run.errors.extend(dpt3.resolve_figure_footnotes(refined, scans, digests))
        # After the furniture sweep, so a removed fragment is never grouped.
        dpt3.group_scanned_text(refined, scans, digests, dims)
        run.errors.extend(dpt3.unhomed_tables(scans, refined, dims))
    # Last, so it judges the document as it will be written: a picture the scan
    # read carries its description by now, and a picture that turned out to be a
    # table has its table beside it.
    run.errors.extend(unread_pictures(refined, dims))

    # After the graft, so each figure record carries its picture_ref and the
    # values it produces can name the picture they belong to.
    values_run: Optional[FigureValueRun] = None
    if cells_pages and value_parser is not None and value_reader is not None:
        figure_pages = sorted({s.page for s in scans if s.figures})
        with tempfile.TemporaryDirectory() as tmp:
            page_images: Dict[int, Path] = {}
            for page_no in figure_pages:
                image_path = Path(tmp) / f"values-p{page_no:04d}.png"
                render_page(source, page_no, CAPTURE_DPI, image_path)
                page_images[page_no] = image_path
            values_run = await read_figure_values(
                scans, cells_pages, page_images, value_parser, value_reader, base
            )
        run.errors.extend(values_run.errors)

    artifacts = [s.response_artifact for s in scans if s.response_artifact]
    with output_sink(output_dir) as out:
        charts_path = out / f"{base}.figures.json"
        charts_path.write_text(run.model_dump_json(indent=2), encoding="utf-8")
        parse_path = out / parse_name
        parse_path.write_text(json.dumps(refined.export_to_dict(), indent=2), encoding="utf-8")
        artifacts.extend([charts_path.name, parse_path.name])
        if tables:
            tables_path = out / f"{base}.scanned-tables.json"
            tables_path.write_text(
                json.dumps([t.model_dump(mode="json") for t in tables], indent=2), encoding="utf-8"
            )
            artifacts.append(tables_path.name)
        if values_run is not None:
            values_path = out / f"{base}.figure-values.json"
            values_path.write_text(values_run.model_dump_json(indent=2), encoding="utf-8")
            artifacts.append(values_path.name)
        if digests:
            # The track's own records: what the projection left out of the graft
            # contract, with each page's resolved and placed footnotes.
            digest_path = out / f"{base}.figure-digest.json"
            digest_path.write_text(
                json.dumps([digests[page].model_dump(mode="json") for page in sorted(digests)], indent=2),
                encoding="utf-8",
            )
            artifacts.append(digest_path.name)

    logger.info(
        "Figure run complete: {} nominated={} scanned={} submitted={} with_figures={} "
        "records={} tables={} credits={}",
        base,
        run.nominated,
        run.scanned,
        run.submitted,
        run.with_figures,
        len(run.figures),
        len(tables),
        run.credits,
    )
    for err in run.errors:
        logger.error("  - {}", err)
    return ChartOutput(
        run=run,
        document=refined,
        tables=tables,
        artifacts=artifacts,
        values=values_run,
        digests=digests,
    )


def run_figures_sync(
    source: Path,
    document: DoclingDocument,
    base: str,
    parse_name: str,
    output_dir: str = "output",
    model: str = SCAN_DEFAULT_MODEL,
    concurrency: int = SCAN_CONCURRENCY,
    pages: str = "nominated",
    llm: Optional[LLMClient] = None,
    inspector: Optional[StatusInspector] = None,
    cell_reader: Optional[CellReader] = None,
    corrector: Optional[FigureCorrector] = None,
    cells_pages: Optional[Dict[int, dict]] = None,
    value_parser: Optional[ScanValueParser] = None,
    value_reader: Optional[LocalValueReader] = None,
) -> ChartOutput:
    """`run_figures` for a synchronous caller."""
    return asyncio.run(
        run_figures(
            source,
            document,
            base,
            parse_name,
            output_dir,
            model,
            concurrency,
            pages,
            llm,
            inspector,
            cell_reader,
            corrector,
            cells_pages,
            value_parser,
            value_reader,
        )
    )


async def _scan_pages(
    source: Path,
    base: str,
    kept: List[NominatedPage],
    model: str,
    output_dir: str,
    concurrency: int,
) -> Tuple[List[PageScan], Dict[int, dpt3.PageDigest], List[str]]:
    """Scan every page in `kept`, writing each raw response before reading it.

    This is where the model selects the path: a dpt-3 model's responses are
    digested by the track and the digests are returned beside the scans. Any
    other model's are read by `page_scan`.

    A page whose scan raises, or whose dpt-3 response cannot be digested, is
    surfaced as an error and gets no record, so the failure is visible rather
    than looking like a page with no chart. A failure in `page_scan` or in the
    dpt-3 projection is not caught and stops the run.
    """
    if not kept:
        return [], {}, []

    semaphore = asyncio.Semaphore(concurrency)
    track = is_v2_model(model)

    async def one(
        nomination: NominatedPage,
    ) -> Tuple[Optional[PageScan], Optional[dpt3.PageDigest], Optional[str]]:
        async with semaphore:
            try:
                response, artifact, reused = await _page_response(
                    source, base, nomination.page, model, output_dir
                )
            except Exception as exc:
                return (
                    None,
                    None,
                    (
                        f"page {nomination.page}: figure scan failed ({type(exc).__name__}: {exc}); "
                        "the page holds a chart the parse did not read"
                    ),
                )
        if track:
            try:
                digest = dpt3.digest_page(response, nomination.page, model, artifact)
            except Exception as exc:
                return (
                    None,
                    None,
                    (
                        f"page {nomination.page}: the scan was billed and its response could not be "
                        f"read ({type(exc).__name__}: {exc}); the raw response is stored as {artifact}"
                    ),
                )
            scan = dpt3.project_scan(digest, nomination.picture_classes, artifact, reused)
            scan.table_refs = list(nomination.table_refs)
            return scan, digest, None
        scan = page_scan(response, nomination.page, model, nomination.picture_classes, artifact, reused)
        scan.table_refs = list(nomination.table_refs)
        return scan, None, None

    results = await asyncio.gather(*(one(n) for n in kept))
    scans = [s for s, _d, _e in results if s is not None]
    digests = {d.page: d for _s, d, _e in results if d is not None}
    errors = [e for _s, _d, e in results if e is not None]
    return scans, digests, errors


async def _page_response(
    source: Path,
    base: str,
    page: int,
    model: str,
    output_dir: str,
) -> Tuple[Dict[str, Any], str, bool]:
    """The raw response for one page, and whether it came from a stored scan.

    The stored response is named the way a single-page ADE run names it, so a
    scan of that page from either path is reused instead of paid for again.
    Stored responses are found only in a local `output_dir`. An `s3://` sink
    starts empty, so there every page is scanned again.

    Reuse is keyed on the model: a stored response is reused only when its
    format matches the requested model's generation, which the file itself
    shows — a chunks list is dpt-2, a structure tree is dpt-3. A mismatched
    file is moved aside to `{base}.p{N}.ade.{generation}.json`, replacing any
    file already there. If a response of the requested generation is waiting
    under `{base}.p{N}.ade.{requested}.json`, it is swapped in and reused.
    Otherwise the page is rescanned. A stored file in neither format raises,
    so a billed page can never pass as blank.
    """
    artifact = f"{base}.p{page}.ade.json"
    requested = "dpt3" if is_v2_model(model) else "dpt2"
    with output_sink(output_dir) as out:
        stored = out / artifact
        if stored.exists():
            response = json.loads(stored.read_text(encoding="utf-8"))
            generation = dpt3.response_generation(response)
            if generation == requested:
                logger.info("Figure scan: reusing the stored response for page {} ({})", page, artifact)
                return response, artifact, True
            if generation is None:
                raise ValueError(
                    f"the stored response {artifact} (model {model}) holds neither a dpt-2 chunks "
                    "list nor a dpt-3 structure tree; move it aside to rescan the page"
                )
            aside = out / f"{base}.p{page}.ade.{generation}.json"
            wanted = out / f"{base}.p{page}.ade.{requested}.json"
            if wanted.exists():
                candidate = json.loads(wanted.read_text(encoding="utf-8"))
                if dpt3.response_generation(candidate) == requested:
                    stored.replace(aside)
                    wanted.replace(stored)
                    logger.info(
                        "Figure scan: page {} stored response is {}; moved aside to {} and "
                        "reusing the {} response already paid for",
                        page,
                        generation,
                        aside.name,
                        requested,
                    )
                    return candidate, artifact, True
            stored.replace(aside)
            logger.info(
                "Figure scan: page {} stored response is {}; moved aside to {} and rescanning " "with {}",
                page,
                generation,
                aside.name,
                model,
            )

    response = await scan_page(source, page, model=model)
    with output_sink(output_dir) as out:
        (out / artifact).write_text(json.dumps(response, indent=1), encoding="utf-8")
    return response, artifact, False


def widen_to_all_pages(document: DoclingDocument, nominated: List[NominatedPage]) -> List[NominatedPage]:
    """Every page of the document, keeping what nominated the nominated ones.

    A page the parse pointed at keeps its picture and table refs; every other
    page joins bare — no refs, just the page — and is scanned whole through
    the same flow.
    """
    covered = {n.page for n in nominated}
    extra = [NominatedPage(page=p) for p in sorted(document.pages) if p not in covered]
    if extra:
        logger.info(
            "Figure run: scanning {} un-nominated page(s) under pages=all: {}",
            len(extra),
            [n.page for n in extra],
        )
    return sorted([*nominated, *extra], key=lambda n: n.page)


def _stored_track_pages(base: str, output_dir: str, nominated: set[int]) -> List[NominatedPage]:
    """Pages beyond nomination whose stored dpt-3 response already sits in a local `output_dir`.

    Reading one costs nothing — the response was paid for by whichever run made
    it, `quber ade --page` included, and it is named the way this workflow
    names its own. Only a response already in the dpt-3 format qualifies: a
    dpt-2 response here would be rescanned at a price, and a page nobody
    nominated must never be billed by the back door.
    """
    pattern = re.compile(rf"{re.escape(base)}\.p(\d+)\.ade\.json")
    found: List[NominatedPage] = []
    with output_sink(output_dir) as out:
        for path in sorted(out.glob(f"{base}.p*.ade.json")):
            match = pattern.fullmatch(path.name)
            if match is None:
                continue
            page = int(match.group(1))
            if page in nominated:
                continue
            try:
                response = json.loads(path.read_text(encoding="utf-8"))
            except ValueError:
                continue
            if dpt3.response_generation(response) == "dpt3":
                found.append(NominatedPage(page=page))
    return found


def _page_dims(document: DoclingDocument) -> Dict[int, Tuple[float, float]]:
    """Each page's width and height in points, as the parse recorded them."""
    return {
        page_no: (page.size.width, page.size.height)
        for page_no, page in document.pages.items()
        if page.size is not None
    }


#: The names the pipeline writes a parse under, newest stage first. A directory
#: given as the parse is read for the first of these it holds, in this order,
#: so a document that has been through fusion is refined rather than the raw
#: parse beside it.
PARSE_NAMES = ("unified.json", "docling.json")


def resolve_parse(parse: str, base: str) -> Path:
    """The parse file to refine, from a JSON file or a directory holding one.

    A file is taken as given. A directory is searched for `<base>.unified.json`
    and then `<base>.docling.json`.
    """
    path = Path(parse)
    if not path.is_dir():
        if not path.exists():
            raise FileNotFoundError(f"no parse to refine at {path}")
        return path
    for name in PARSE_NAMES:
        candidate = path / f"{base}.{name}"
        if candidate.exists():
            return candidate
    wanted = " or ".join(f"{base}.{n}" for n in PARSE_NAMES)
    raise FileNotFoundError(f"no parse to refine in {path}: expected {wanted}")


def load_document(parse: str, base: str) -> Tuple[DoclingDocument, str]:
    """The parse to refine and the filename it was read from."""
    path = resolve_parse(parse, base)
    return DoclingDocument.load_from_json(path), path.name


def write_marker(output_dir: str, base: str, payload: Dict[str, Any]) -> None:
    """Write the figure run's completion marker, strictly after its artifacts."""
    write_completion_marker(output_dir, f"{base}.figure.complete.json", payload)