from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING, Any

import aicd.tools.viz_assets as viz_assets
from aicd.tools import html_screenshot as html_screenshot_mod
from aicd.kit import chart_spec
from aicd.kit.validate_viz import validate_html_viz_embed

if TYPE_CHECKING:
    from aicd.agent.tool_router import ToolRouter


def _ai_output_info(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    return {
        "ok": True,
        "workspace": str(router._cwd),
        "ai_output_dir": str(router._ai_output),
        "ai_tmp_dir": str(router._ai_tmp),
        "note": "Write scratch/intermediate analysis under ai_tmp_dir; clean up when done. "
        "Long-lived corpus: **ai_output_dir/output/library/<doc_slug>/** — keep **INDEX.md** with **latest:** pointer + **## Revisions**; subdirs **apis/**, **diagrams/**, **outlines/**, **markdown/**, optional **chapters/chNN_*/**. "
        "Doc iterations: new **markdown/*_vNNN*.md** (bump NNN), update INDEX; see system DOC_AUTHORING_ITERATION_DIRECTIVES. "
        "Durable artifacts go under ai_output_dir (not inside tmp unless explicitly temporary). "
        "**script_run** archives executed scripts under **ai_output_dir/bak/** (timestamp + **related_doc** in filename). "
        "Offline **viz_export_html** / **viz_data_chart** outputs (**.html** + sibling **res/**) are durable preview assets—do not delete; prefer **output/library/.../diagrams/** or deliverables/.",
    }


def _ai_ensure_layout(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    router._ai_output.mkdir(parents=True, exist_ok=True)
    router._ai_tmp.mkdir(parents=True, exist_ok=True)
    (router._ai_output / "bak").mkdir(parents=True, exist_ok=True)
    return {
        "ok": True,
        "ai_output_dir": str(router._ai_output),
        "ai_tmp_dir": str(router._ai_tmp),
        "script_backup_dir": str((router._ai_output / "bak").resolve()),
        "created": True,
    }


def _ai_workspace_set(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    fn = router._ai_workspace_setter
    if fn is None:
        return {
            "ok": False,
            "error": "ai_workspace_set is only available in the main chat session; use TUI /work <path>.",
        }
    raw = args.get("path")
    if raw is None or (isinstance(raw, str) and not str(raw).strip()):
        ok, msg = fn(None)
    else:
        ok, msg = fn(str(raw).strip())
    if not ok:
        return {"ok": False, "error": msg}
    return {
        "ok": True,
        "workspace": msg,
        "note": "fs_* relative paths and default run_command cwd use this directory.",
    }


def _viz_list_bundled_assets(
    router: ToolRouter, args: dict[str, Any]
) -> dict[str, Any]:
    return {
        "ok": True,
        "bundled_dir": str(viz_assets.bundled_res_dir()),
        "files": viz_assets.list_bundled_viz_files(),
    }


def _viz_copy_resources(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    dest = router._resolve_path(args["target_directory"])
    if not dest.is_dir():
        return {"ok": False, "error": "target_directory must be an existing directory"}
    return viz_assets.copy_viz_resources(dest)


def _viz_export_html(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    out = router._resolve_path(args["html_path"])
    out.parent.mkdir(parents=True, exist_ok=True)
    title = str(args.get("title") or "Analysis")
    raw_body = args.get("body_html")
    body = (raw_body if raw_body is not None else "").strip()
    if not body:
        body = viz_assets.demo_body_snippet()
        boot = True
    else:
        body = str(raw_body)
        boot = bool(args.get("include_demo_chart_vis_bootstrap", False))
    html = viz_assets.build_viz_html_document(
        title,
        body,
        extra_head=str(args.get("extra_head") or ""),
        extra_scripts=str(args.get("extra_scripts") or ""),
        include_default_chart_vis_init=boot,
    )
    html = viz_assets.reorder_body_scripts_before_vis_usage(html)
    out.write_text(html, encoding="utf-8")
    copied: dict[str, Any] = {"skipped": True}
    if bool(args.get("copy_resources", True)):
        copied = viz_assets.copy_viz_resources(out.parent)
    embed_v = validate_html_viz_embed(html)
    ret: dict[str, Any] = {
        "ok": True,
        "html_path": str(out),
        "copy_resources": copied,
    }
    if embed_v.get("issues"):
        ret["html_embed_validation"] = embed_v
    return ret


def _viz_data_chart(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    out = router._resolve_path(args["html_path"])
    out.parent.mkdir(parents=True, exist_ok=True)
    title = str(args.get("title") or "Chart")
    chart_type = str(args.get("chart_type") or "bar")
    labels = list(args.get("labels") or [])
    values = list(args.get("values") or [])
    dataset_label = str(args.get("dataset_label") or "")
    html = chart_spec.build_html_with_chart(
        title,
        chart_type,
        labels,
        values,
        dataset_label=dataset_label or "Series",
    )
    out.write_text(html, encoding="utf-8")
    copied: dict[str, Any] = {"skipped": True}
    if bool(args.get("copy_resources", True)):
        copied = viz_assets.copy_viz_resources(out.parent)
    embed_v = validate_html_viz_embed(html)
    ret_dc: dict[str, Any] = {
        "ok": True,
        "html_path": str(out),
        "copy_resources": copied,
        "ascii_preview": chart_spec.mermaid_bar_from_counts(
            [str(x) for x in labels],
            [int(round(float(v))) for v in values],
        )
        if len(labels) == len(values)
        else "",
    }
    if embed_v.get("issues"):
        ret_dc["html_embed_validation"] = embed_v
    return ret_dc


async def _viz_html_to_png(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    """离线 HTML → PNG；不删除源 HTML（与 viz_export_html / viz_data_chart 配套）。"""
    hp = router._resolve_path(str(args["html_path"]))
    op = router._resolve_path(str(args["output_png_path"]))
    sel = args.get("selector")
    selector = str(sel).strip() if sel is not None and str(sel).strip() else None
    full_page = bool(args.get("full_page", True))
    dsf = float(args.get("device_scale_factor", 2))
    wait_ms = int(args.get("wait_after_load_ms", 1200))
    timeout_ms = int(args.get("timeout_ms", 120_000))
    policy = router._policy.check_allowed

    def _sync() -> dict[str, Any]:
        return html_screenshot_mod.render_html_file_to_png(
            hp,
            op,
            policy_check=policy,
            selector=selector,
            full_page=full_page,
            device_scale_factor=dsf,
            wait_after_load_ms=wait_ms,
            timeout_ms=timeout_ms,
        )

    return await asyncio.to_thread(_sync)


HANDLERS: dict[str, Any] = {
    "ai_output_info": _ai_output_info,
    "ai_ensure_layout": _ai_ensure_layout,
    "ai_workspace_set": _ai_workspace_set,
    "viz_list_bundled_assets": _viz_list_bundled_assets,
    "viz_copy_resources": _viz_copy_resources,
    "viz_export_html": _viz_export_html,
    "viz_data_chart": _viz_data_chart,
    "viz_html_to_png": _viz_html_to_png,
}
