"""
Visualize Capability
====================

Unified visualization capability. AnalysisAgent picks one of six render
types — svg / chartjs / mermaid / html (text-emitting, three-stage pipeline)
or manim_video / manim_image (Manim subprocess pipeline). The result
envelope carries ``render_type`` as the discriminator so the frontend can
delegate to the right viewer.
"""

from __future__ import annotations

import logging
from typing import Any

from deeptutor.agents._shared.capability_result import emit_capability_result
from deeptutor.core.agentic.usage import UsageTracker
from deeptutor.core.capability_protocol import BaseCapability, CapabilityManifest
from deeptutor.core.context import UnifiedContext
from deeptutor.core.stream_bus import StreamBus
from deeptutor.core.trace import merge_trace_metadata
from deeptutor.i18n import StatusI18n
from deeptutor.runtime.request_contracts import (
    VisualizeRequestConfig,
    get_capability_request_schema,
    validate_visualize_request_config,
)

logger = logging.getLogger(__name__)

# Stages exposed in the manifest. The first three cover the text-emitting
# path (svg/chartjs/mermaid/html); the rest cover the manim subprocess
# path. A given turn only streams a subset of these.
_VISUALIZE_STAGES = [
    "analyzing",
    "generating",
    "reviewing",
    "concept_analysis",
    "concept_design",
    "code_generation",
    "code_retry",
    "summary",
    "render_output",
]

_MANIM_RENDER_TYPES = {"manim_video", "manim_image"}


class VisualizeCapability(BaseCapability):
    manifest = CapabilityManifest(
        name="visualize",
        description=(
            "Generate SVG, Chart.js, Mermaid, interactive HTML, or Manim "
            "animation/storyboard visualizations."
        ),
        stages=_VISUALIZE_STAGES,
        tools_used=[],
        cli_aliases=["visualize", "viz"],
        request_schema=get_capability_request_schema("visualize"),
    )

    async def run(self, context: UnifiedContext, stream: StreamBus) -> None:
        from deeptutor.agents.visualize.models import ReviewResult
        from deeptutor.agents.visualize.pipeline import VisualizePipeline
        from deeptutor.agents.visualize.utils import (
            build_fallback_html,
            validate_visualization,
        )
        from deeptutor.services.llm.config import get_llm_config

        request_config = validate_visualize_request_config(context.config_overrides)
        render_mode = request_config.render_mode
        i18n = StatusI18n(self.name, context.language, module="visualize")

        llm_config_for_usage = get_llm_config()
        usage = UsageTracker(model=getattr(llm_config_for_usage, "model", None))

        llm_config = get_llm_config()
        history_context = str(context.metadata.get("conversation_context_text", "") or "").strip()

        pipeline = VisualizePipeline(
            api_key=llm_config.api_key,
            base_url=llm_config.base_url,
            api_version=llm_config.api_version,
            language=context.language,
            trace_callback=self._build_trace_bridge(stream, i18n=i18n),
        )

        # Stage 1: Analyze (routing decision)
        async with stream.stage("analyzing", source=self.name):
            await stream.thinking(
                i18n.t("analyzing", "Analyzing visualization requirements..."),
                source=self.name,
                stage="analyzing",
            )
            analysis = await pipeline.run_analysis(
                user_input=context.user_message,
                history_context=history_context,
                render_mode=render_mode,
                attachments=context.attachments,
            )
            await stream.progress(
                message=i18n.t(
                    "render_type_detected",
                    f"Render type: {analysis.render_type} — {analysis.description}",
                    render_type=analysis.render_type,
                    description=analysis.description,
                ),
                source=self.name,
                stage="analyzing",
            )

        # Branch: manim path takes over completely after the analysis stage,
        # using its own multi-agent pipeline + Manim subprocess.
        if analysis.render_type in _MANIM_RENDER_TYPES:
            await self._run_manim_path(
                context=context,
                stream=stream,
                render_type=analysis.render_type,
                visualize_config=request_config,
                history_context=history_context,
                usage=usage,
                i18n=i18n,
            )
            return

        # Stage 2: Generate code
        async with stream.stage("generating", source=self.name):
            await stream.thinking(
                i18n.t("generating", "Generating visualization code..."),
                source=self.name,
                stage="generating",
            )
            code = await pipeline.run_code_generation(
                user_input=context.user_message,
                history_context=history_context,
                analysis=analysis,
            )
            await stream.progress(
                message=i18n.t("code_generated", "Code generated."),
                source=self.name,
                stage="generating",
            )

        # Stage 3: Validate locally; repair only on failure.
        #
        # The old generic LLM review is replaced by a deterministic, zero-cost
        # local check (well-formed XML / strict-JSON / mermaid lint / HTML
        # sanity). When it passes we ship the draft as-is — saving a whole
        # serial LLM call. When it fails we spend one *targeted* repair call
        # driven by the concrete error, not an open-ended re-judgement.
        async with stream.stage("reviewing", source=self.name):
            ok, validation_error = validate_visualization(code, analysis.render_type)
            if ok:
                final_code = code
                review = ReviewResult(
                    optimized_code=final_code,
                    changed=False,
                    review_notes="Passed local validation.",
                )
                await stream.progress(
                    message=i18n.t(
                        "validation_passed",
                        "Looks good — passed local checks.",
                    ),
                    source=self.name,
                    stage="reviewing",
                )
            elif analysis.render_type == "html":
                # html documents are 8-16k tokens; we don't run them through
                # the repair loop — fall back to a minimal renderable template.
                final_code = build_fallback_html(
                    title=analysis.description or "Visualization",
                    summary=analysis.data_description,
                    note="The model did not return a renderable HTML document.",
                )
                review = ReviewResult(
                    optimized_code=final_code,
                    changed=True,
                    review_notes=f"Used fallback HTML template ({validation_error}).",
                )
                await stream.progress(
                    message=i18n.t(
                        "html_invalid_fallback",
                        "HTML did not validate; using fallback template.",
                    ),
                    source=self.name,
                    stage="reviewing",
                )
            else:
                await stream.thinking(
                    i18n.t("repairing", "Fixing a validation issue..."),
                    source=self.name,
                    stage="reviewing",
                )
                try:
                    review = await pipeline.run_repair(
                        user_input=context.user_message,
                        analysis=analysis,
                        code=code,
                        error=validation_error,
                    )
                except Exception as exc:
                    # Repair wraps code inside a JSON string field; large/complex
                    # SVGs can trip JSON-mode escaping. Fall back to the draft so
                    # the user still gets a rendered result.
                    logger.warning("Visualize repair failed (%s); using unvalidated draft.", exc)
                    review = ReviewResult(
                        optimized_code=code,
                        changed=False,
                        review_notes=f"Repair skipped due to error: {exc}",
                    )
                    final_code = code
                    await stream.progress(
                        message=i18n.t(
                            "repair_skipped_error",
                            "Repair skipped — using draft as-is.",
                        ),
                        source=self.name,
                        stage="reviewing",
                    )
                else:
                    final_code = review.optimized_code or code
                    repaired_ok, repaired_error = validate_visualization(
                        final_code, analysis.render_type
                    )
                    if repaired_ok:
                        await stream.progress(
                            message=i18n.t(
                                "code_repaired",
                                f"Fixed: {review.review_notes}",
                                notes=review.review_notes,
                            ),
                            source=self.name,
                            stage="reviewing",
                        )
                    else:
                        await stream.progress(
                            message=i18n.t(
                                "repair_incomplete",
                                f"Repair attempted; residual issue: {repaired_error}",
                                error=repaired_error,
                            ),
                            source=self.name,
                            stage="reviewing",
                        )

        # Emit final content as a fenced code block for the chat area
        if analysis.render_type == "svg":
            lang_tag = "svg"
        elif analysis.render_type == "mermaid":
            lang_tag = "mermaid"
        elif analysis.render_type == "html":
            lang_tag = "html"
        else:
            lang_tag = "javascript"
        content_md = f"```{lang_tag}\n{final_code}\n```"
        await stream.content(content_md, source=self.name, stage="reviewing")

        # Structured result for the frontend viewer
        await emit_capability_result(
            stream,
            {
                "response": content_md,
                "render_type": analysis.render_type,
                "code": {
                    "language": lang_tag,
                    "content": final_code,
                },
                "analysis": analysis.model_dump(),
                "review": review.model_dump(),
            },
            source=self.name,
            usage=usage,
        )

    async def _run_manim_path(
        self,
        *,
        context: UnifiedContext,
        stream: StreamBus,
        render_type: str,
        visualize_config: VisualizeRequestConfig,
        history_context: str,
        usage: UsageTracker | None = None,
        i18n: StatusI18n | None = None,
    ) -> None:
        """
        Manim sub-pipeline. Mirrors ``MathAnimatorCapability.run`` but emits
        the final result with ``render_type`` as the discriminator so the
        unified frontend dispatcher can route to ``MathAnimatorViewer``.
        """
        import importlib.util
        import time

        if importlib.util.find_spec("manim") is None:
            raise RuntimeError(
                "Manim rendering requires optional dependencies. "
                "Install with `pip install 'deeptutor[math-animator]'` "
                "or `pip install -r requirements/math-animator.txt`."
            )

        from deeptutor.agents.math_animator.pipeline import MathAnimatorPipeline
        from deeptutor.agents.math_animator.request_config import MathAnimatorRequestConfig
        from deeptutor.core.trace import build_trace_metadata, new_call_id
        from deeptutor.services.llm.config import get_llm_config

        if i18n is None:
            i18n = StatusI18n(self.name, context.language, module="visualize")
        output_mode = "image" if render_type == "manim_image" else "video"
        request_config = MathAnimatorRequestConfig(
            output_mode=output_mode,  # type: ignore[arg-type]
            quality=visualize_config.quality,
            style_hint=visualize_config.style_hint,
        )

        llm_config = get_llm_config()
        pipeline = MathAnimatorPipeline(
            api_key=llm_config.api_key,
            base_url=llm_config.base_url,
            api_version=llm_config.api_version,
            language=context.language,
            trace_callback=self._build_trace_bridge(stream, i18n=i18n),
        )

        timings: dict[str, float] = {}
        turn_id = str(
            context.metadata.get("turn_id", "") or context.session_id or "visualize-manim"
        )
        render_call_meta = build_trace_metadata(
            call_id=new_call_id("manim-render"),
            phase="render_output",
            label="Render output",
            call_kind="math_render_output",
            trace_role="render",
            trace_kind="progress",
            output_mode=request_config.output_mode,
            quality=request_config.quality,
        )

        stage_start = time.perf_counter()
        async with stream.stage("concept_analysis", source=self.name):
            analysis = await pipeline.run_analysis(
                user_input=context.user_message,
                history_context=history_context,
                request_config=request_config,
                attachments=context.attachments,
            )
        timings["concept_analysis"] = round(time.perf_counter() - stage_start, 3)

        stage_start = time.perf_counter()
        async with stream.stage("concept_design", source=self.name):
            design = await pipeline.run_design(
                user_input=context.user_message,
                request_config=request_config,
                analysis=analysis,
            )
        timings["concept_design"] = round(time.perf_counter() - stage_start, 3)

        stage_start = time.perf_counter()
        async with stream.stage("code_generation", source=self.name):
            generated = await pipeline.run_code_generation(
                user_input=context.user_message,
                request_config=request_config,
                analysis=analysis,
                design=design,
            )
            await stream.progress(
                message=i18n.t("manim_code_prepared", "Manim code prepared."),
                source=self.name,
                stage="code_generation",
            )
        timings["code_generation"] = round(time.perf_counter() - stage_start, 3)

        async def _on_retry(retry_attempt) -> None:
            await stream.progress(
                message=i18n.t(
                    "manim_retry",
                    f"Retry {retry_attempt.attempt}: {retry_attempt.error}",
                    attempt=retry_attempt.attempt,
                    error=retry_attempt.error,
                ),
                source=self.name,
                stage="code_retry",
                metadata={**render_call_meta, "trace_layer": "raw"},
            )

        async def _on_render_progress(message: str, raw: bool) -> None:
            await stream.progress(
                message=message,
                source=self.name,
                stage="render_output",
                metadata={
                    **render_call_meta,
                    "trace_layer": "raw" if raw else "summary",
                },
            )

        async def _on_retry_status(message: str) -> None:
            await stream.progress(
                message=message,
                source=self.name,
                stage="code_retry",
                metadata={"trace_layer": "summary"},
            )

        stage_start = time.perf_counter()
        async with stream.stage("code_retry", source=self.name):
            await stream.progress(
                message=i18n.t(
                    "manim_rendering",
                    (
                        f"Rendering {request_config.output_mode} "
                        f"with quality={request_config.quality}."
                    ),
                    mode=request_config.output_mode,
                    quality=request_config.quality,
                ),
                source=self.name,
                stage="code_retry",
                metadata={**render_call_meta, "call_state": "running"},
            )
            final_code, render_result = await pipeline.run_render(
                turn_id=turn_id,
                user_input=context.user_message,
                request_config=request_config,
                initial_code=generated.code,
                on_retry=_on_retry,
                on_render_progress=_on_render_progress,
                on_retry_status=_on_retry_status,
            )
        timings["code_retry"] = round(time.perf_counter() - stage_start, 3)

        stage_start = time.perf_counter()
        async with stream.stage("summary", source=self.name):
            summary = await pipeline.run_summary(
                user_input=context.user_message,
                request_config=request_config,
                analysis=analysis,
                design=design,
                render_result=render_result,
            )
            if summary.summary_text:
                await stream.content(summary.summary_text, source=self.name, stage="summary")
        timings["summary"] = round(time.perf_counter() - stage_start, 3)

        async with stream.stage("render_output", source=self.name):
            artifact_count = len(render_result.artifacts)
            artifact_key = "manim_artifacts_one" if artifact_count == 1 else "manim_artifacts_many"
            await stream.progress(
                message=i18n.t(
                    artifact_key,
                    (
                        f"Prepared {artifact_count} "
                        f"{'artifact' if artifact_count == 1 else 'artifacts'}."
                    ),
                    count=artifact_count,
                ),
                source=self.name,
                stage="render_output",
                metadata={**render_call_meta, "call_state": "complete"},
            )
        timings["render_output"] = 0.0
        visual_review = getattr(render_result, "visual_review", None)

        await emit_capability_result(
            stream,
            {
                "response": summary.summary_text,
                "render_type": render_type,
                "summary": summary.model_dump(),
                "code": {
                    "language": "python",
                    "content": final_code,
                },
                "output_mode": request_config.output_mode,
                "artifacts": [artifact.model_dump() for artifact in render_result.artifacts],
                "timings": timings,
                "render": {
                    "quality": request_config.quality,
                    "retry_attempts": render_result.retry_attempts,
                    "retry_history": [item.model_dump() for item in render_result.retry_history],
                    "source_code_path": render_result.source_code_path,
                    "visual_review": visual_review.model_dump() if visual_review else None,
                },
                "analysis": analysis.model_dump(),
                "design": design.model_dump(),
            },
            source=self.name,
            usage=usage,
        )

    def _build_trace_bridge(self, stream: StreamBus, i18n: StatusI18n | None = None):
        async def _trace_bridge(update: dict[str, Any]) -> None:
            event = str(update.get("event", "") or "")
            stage = str(update.get("phase") or update.get("stage") or "analyzing")
            base_metadata = {
                key: value
                for key, value in update.items()
                if key
                not in {"event", "state", "response", "chunk", "result", "tool_name", "tool_args"}
            }

            if event != "llm_call":
                return

            state = str(update.get("state", "running"))
            label = str(base_metadata.get("label", "") or stage.replace("_", " ").title())
            if state == "running":
                await stream.progress(
                    message=label,
                    source=self.name,
                    stage=stage,
                    metadata=merge_trace_metadata(
                        base_metadata,
                        {"trace_kind": "call_status", "call_state": "running"},
                    ),
                )
                return
            if state == "streaming":
                chunk = str(update.get("chunk", "") or "")
                if chunk:
                    await stream.thinking(
                        chunk,
                        source=self.name,
                        stage=stage,
                        metadata=merge_trace_metadata(
                            base_metadata,
                            {"trace_kind": "llm_chunk"},
                        ),
                    )
                return
            if state == "complete":
                was_streaming = update.get("streaming", False)
                if not was_streaming:
                    response = str(update.get("response", "") or "")
                    if response:
                        await stream.thinking(
                            response,
                            source=self.name,
                            stage=stage,
                            metadata=merge_trace_metadata(
                                base_metadata,
                                {"trace_kind": "llm_output"},
                            ),
                        )
                await stream.progress(
                    message=label,
                    source=self.name,
                    stage=stage,
                    metadata=merge_trace_metadata(
                        base_metadata,
                        {"trace_kind": "call_status", "call_state": "complete"},
                    ),
                )
                return
            if state == "error":
                fallback = (
                    i18n.t("llm_call_failed", "LLM call failed.")
                    if i18n is not None
                    else "LLM call failed."
                )
                await stream.error(
                    str(update.get("response", "") or fallback),
                    source=self.name,
                    stage=stage,
                    metadata=merge_trace_metadata(
                        base_metadata,
                        {"trace_kind": "call_status", "call_state": "error"},
                    ),
                )

        return _trace_bridge