from __future__ import annotations

from pathlib import Path

import gradio as gr

from models.model_catalog import ModelInfo
from training.export import (
    QUANTIZATION_CHOICES,
    build_export_plan,
    list_exported_files,
)
from ui.progress import CLICK_PROGRESS


def export_download_paths(plan: dict | None, exported_rows: list[list[str]]) -> list[str]:
    """Return existing exported files for the Gradio download output."""
    existing_paths = [
        Path(row[0])
        for row in exported_rows
        if row and Path(row[0]).is_file()
    ]
    if not plan:
        return [str(path) for path in existing_paths]

    planned_paths = _planned_export_paths(plan)
    planned_keys = {_path_key(path) for path in planned_paths}
    planned_existing = [
        path for path in existing_paths if _path_key(path) in planned_keys
    ]
    other_existing = [
        path for path in existing_paths if _path_key(path) not in planned_keys
    ]
    return [str(path) for path in [*planned_existing, *other_existing]]


def _planned_export_paths(plan: dict) -> list[Path]:
    output_dir = Path(str(plan.get("output_dir") or ""))
    paths: list[Path] = []

    official_file = str(plan.get("official_gguf_file") or "")
    if official_file:
        paths.append(output_dir / official_file)

    convert_command = plan.get("convert_command") or []
    if "--outfile" in convert_command:
        outfile_index = convert_command.index("--outfile") + 1
        if outfile_index < len(convert_command):
            paths.append(Path(str(convert_command[outfile_index])))

    quantize_command = plan.get("quantize_command") or []
    if len(quantize_command) >= 3:
        paths.append(Path(str(quantize_command[2])))

    return paths


def _path_key(path: Path) -> str:
    return str(path.resolve())


def build_export_tab(catalog: dict[str, ModelInfo]) -> None:
    model_id = gr.Dropdown(list(catalog), value=next(iter(catalog)), label="Model")
    quant = gr.Dropdown(QUANTIZATION_CHOICES, value="Q4_K_M", label="Quantization")
    output_dir = gr.Textbox(label="Output directory", value="exports")
    run = gr.Button("Prepare export plan", variant="primary")
    refresh = gr.Button("Refresh exported files")
    output = gr.JSON(label="Export plan")
    files = gr.Dataframe(headers=["path", "bytes"], label="Exported files", interactive=False)
    downloads = gr.File(
        label="Download exported files",
        file_count="multiple",
        interactive=False,
    )

    def plan_export(
        selected: str,
        selected_quant: str,
        directory: str,
    ) -> tuple[dict, list[list[str]], list[str]]:
        model = catalog[selected]
        plan = build_export_plan(model, selected_quant, directory)
        exported_rows = list_exported_files(directory)
        plan_data = plan.as_dict()
        return plan_data, exported_rows, export_download_paths(plan_data, exported_rows)

    def refresh_exports(directory: str) -> tuple[list[list[str]], list[str]]:
        exported_rows = list_exported_files(directory)
        return exported_rows, export_download_paths(None, exported_rows)

    run.click(
        plan_export,
        [model_id, quant, output_dir],
        [output, files, downloads],
        show_progress=CLICK_PROGRESS,
    )
    refresh.click(
        refresh_exports,
        output_dir,
        [files, downloads],
        show_progress=CLICK_PROGRESS,
    )