"""Public API for the Yagra package."""

from __future__ import annotations

import argparse
import json
import sys
from collections.abc import Mapping
from os import PathLike
from pathlib import Path
from typing import Any

from langchain_core.runnables.config import RunnableConfig
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.graph.state import CompiledStateGraph

from yagra.adapters.inbound import create_workflow_studio_server
from yagra.adapters.outbound import InMemoryNodeRegistry
from yagra.application.services.template_initializer import (
    FileAlreadyExistsError,
    TemplateNotFoundError,
    initialize_from_template,
    list_templates,
    list_templates_with_info,
)
from yagra.application.use_cases import (
    WorkflowValidationIssue,
    WorkflowValidationReport,
    build_from_workflow_path,
    format_validation_report,
    render_workflow_visualization_html,
    validate_workflow_for_ui,
    validate_workflow_payload_for_ui,
)
from yagra.domain.entities import GraphSpec
from yagra.ports.outbound import NodeHandler, NodeRegistryPort


class Yagra:
    """Wrapper for executing LangGraph built from workflow YAML."""

    def __init__(self, compiled_graph: CompiledStateGraph) -> None:
        """Initializes with a compiled graph.

        Args:
            compiled_graph: Compiled graph generated by `build_state_graph`.
        """
        self._compiled_graph = compiled_graph

    @classmethod
    def from_workflow(
        cls,
        workflow_path: str | PathLike[str],
        registry: NodeRegistryPort | Mapping[str, NodeHandler],
        bundle_root: str | PathLike[str] | None = None,
        state_schema: Any = None,
        checkpointer: BaseCheckpointSaver | None = None,
    ) -> Yagra:
        """Creates a `Yagra` instance from a workflow file.

        If `interrupt_before` / `interrupt_after` are defined in the workflow YAML,
        a `checkpointer` must be specified to enable HITL (Human-in-the-Loop) functionality.
        If `checkpointer` is None, the interrupt configuration is ignored.

        Args:
            workflow_path: Path to the entry `workflow.yaml`.
            registry: Registry implementation resolving handler names to callables, or handler mapping.
            bundle_root: Base directory for resolving split references. Defaults to workflow parent directory if not specified.
            state_schema: LangGraph state schema. If None, automatically resolved
                from the workflow's state_schema definition or defaults to dict.
            checkpointer: LangGraph checkpointer. Required when using `interrupt_before` / `interrupt_after`.

        Returns:
            `Yagra` instance containing the compiled graph.
        """
        normalized_registry = _normalize_registry(registry)
        compiled = build_from_workflow_path(
            workflow_path=workflow_path,
            registry=normalized_registry,
            bundle_root=bundle_root,
            state_schema=state_schema,
            checkpointer=checkpointer,
        )
        return cls(compiled)

    @property
    def compiled_graph(self) -> CompiledStateGraph:
        """Returns the held compiled graph."""
        return self._compiled_graph

    def invoke(
        self,
        state: Mapping[str, Any],
        thread_id: str | None = None,
    ) -> dict[str, Any]:
        """Executes the graph with an initial state.

        When using HITL (Human-in-the-Loop), specify `thread_id`.
        If interrupted by `interrupt_before` / `interrupt_after`, the return value is the state just before interruption.
        Call `resume()` to resume.

        Args:
            state: State dictionary at execution start.
            thread_id: Thread ID used for checkpoint management.
                Specify a unique string when using HITL.
                If None, state is not saved even if a checkpointer is configured.

        Returns:
            State dictionary after execution (or at interruption point).

        Raises:
            TypeError: If the graph execution result is not dict-compatible.
        """
        config: RunnableConfig = (
            {"configurable": {"thread_id": thread_id}} if thread_id is not None else {}
        )
        result = self._compiled_graph.invoke(dict(state), config=config)
        if not isinstance(result, Mapping):
            raise TypeError("compiled graph returned non-mapping result")
        return dict(result)

    def resume(
        self,
        update: Mapping[str, Any] | None = None,
        thread_id: str | None = None,
    ) -> dict[str, Any]:
        """Resumes a graph that was interrupted.

        Resumes a workflow that was interrupted by `invoke()`.
        You can modify the state with `update` before resuming (e.g., injecting human review results).

        Args:
            update: Update dict to merge into state before resuming. If None, resumes without modifying state.
            thread_id: Thread ID specified at `invoke()` time.

        Returns:
            State dictionary after resuming execution.

        Raises:
            TypeError: If the graph execution result is not dict-compatible.
            ValueError: If thread_id is not provided when checkpointer requires it.
        """
        config: RunnableConfig = (
            {"configurable": {"thread_id": thread_id}} if thread_id is not None else {}
        )
        if update is not None:
            self._compiled_graph.update_state(config, dict(update))
        result = self._compiled_graph.invoke(None, config=config)
        if not isinstance(result, Mapping):
            raise TypeError("compiled graph returned non-mapping result")
        return dict(result)


def main() -> None:
    """CLI entry point for Yagra."""
    parser = _build_cli_parser()
    args = parser.parse_args()

    if args.command == "init":
        exit_code = _run_init_command(args)
        raise SystemExit(exit_code)
    if args.command == "schema":
        exit_code = _run_schema_command(args)
        raise SystemExit(exit_code)
    if args.command == "validate":
        exit_code = _run_validate_command(args)
        raise SystemExit(exit_code)
    if args.command == "visualize":
        exit_code = _run_visualize_command(args)
        raise SystemExit(exit_code)
    if args.command == "studio":
        exit_code = _run_studio_command(args)
        raise SystemExit(exit_code)
    if args.command == "handlers":
        exit_code = _run_handlers_command(args)
        raise SystemExit(exit_code)
    if args.command == "explain":
        exit_code = _run_explain_command(args)
        raise SystemExit(exit_code)
    if args.command == "mcp":
        exit_code = _run_mcp_command(args)
        raise SystemExit(exit_code)
    raise SystemExit(2)


def _normalize_registry(registry: NodeRegistryPort | Mapping[str, NodeHandler]) -> NodeRegistryPort:
    """Normalizes the registry argument to a `NodeRegistryPort` implementation.

    Args:
        registry: Registry implementation or handler mapping.

    Returns:
        `NodeRegistryPort` implementation.

    Raises:
        TypeError: If an unsupported type is passed.
    """
    if isinstance(registry, NodeRegistryPort):
        return registry
    if isinstance(registry, Mapping):
        return InMemoryNodeRegistry(registry)
    raise TypeError("registry must be NodeRegistryPort or mapping[str, NodeHandler]")


def _run_init_command(args: argparse.Namespace) -> int:
    """Executes the `init` subcommand.

    Args:
        args: Parsed CLI arguments.

    Returns:
        Exit code. 0 on success.
    """
    # Show template list
    if args.list:
        names = list_templates()
        if not names:
            print("No templates available.")
            return 1
        infos = {info.name: info for info in list_templates_with_info()}
        print("Available templates:")
        for name in names:
            info = infos.get(name)
            if info and info.use_case:
                print(f"  - {name}  [{info.use_case}]")
            else:
                print(f"  - {name}")
            if info and info.description:
                print(f"      {info.description}")
        return 0

    # Error if no template name is specified
    if args.template is None:
        print(
            "Error: specify either --template or --list.",
            file=sys.stderr,
        )
        return 2

    output_dir = Path(args.output).expanduser().resolve()

    try:
        initialize_from_template(
            template_name=args.template,
            output_dir=output_dir,
            force=args.force,
        )
    except TemplateNotFoundError as e:
        print(str(e), file=sys.stderr)
        return 1
    except FileAlreadyExistsError as e:
        print(str(e), file=sys.stderr)
        return 1

    print(f"Initialized from template '{args.template}': {output_dir}")

    # Validate the generated workflow
    workflow_path = output_dir / "workflow.yaml"
    if workflow_path.exists():
        print(f"\nValidating workflow: {workflow_path}")
        report = validate_workflow_for_ui(
            workflow_path=str(workflow_path),
            bundle_root=None,
        )
        if report.is_valid:
            print("✓ workflow is valid.")
        else:
            print("✗ workflow has issues:", file=sys.stderr)
            print(format_validation_report(report), file=sys.stderr)
            return 1

    return 0


def _build_cli_parser() -> argparse.ArgumentParser:
    """Builds the CLI parser.

    Returns:
        Instance of `argparse.ArgumentParser`.
    """
    parser = argparse.ArgumentParser(prog="yagra")
    subparsers = parser.add_subparsers(dest="command", required=True)

    init = subparsers.add_parser(
        "init",
        help="Initialize a workflow from a template",
    )
    init.add_argument(
        "--template",
        default=None,
        help="Template name to use (e.g. branch, loop, rag)",
    )
    init.add_argument(
        "--list",
        action="store_true",
        help="List available templates",
    )
    init.add_argument(
        "--force",
        action="store_true",
        help="Overwrite existing files",
    )
    init.add_argument(
        "--output",
        default=".",
        help="Output directory (default: current directory)",
    )

    schema = subparsers.add_parser(
        "schema",
        help="Output the JSON Schema for workflow YAML",
    )
    schema.add_argument(
        "--output",
        default=None,
        help="Output file path (stdout if not specified)",
    )

    validate = subparsers.add_parser(
        "validate",
        help="Validate workflow YAML",
    )
    validate.add_argument(
        "--workflow",
        required=True,
        help="Path to the workflow YAML to validate",
    )
    validate.add_argument(
        "--bundle-root",
        default=None,
        help="Base directory for resolving split references",
    )
    validate.add_argument(
        "--format",
        choices=["text", "json"],
        default="text",
        help="Output format (default: text)",
    )

    visualize = subparsers.add_parser(
        "visualize",
        help="Visualize a workflow as read-only HTML",
    )
    visualize.add_argument(
        "--workflow",
        required=True,
        help="Path to the workflow YAML to visualize",
    )
    visualize.add_argument(
        "--bundle-root",
        default=None,
        help="Base directory for resolving split references",
    )
    visualize.add_argument(
        "--output",
        default="workflow-visualization.html",
        help="Path for the generated HTML file",
    )
    visualize.add_argument(
        "--title",
        default=None,
        help="Title for the visualization page",
    )

    studio = subparsers.add_parser(
        "studio",
        help="Start the local WebUI/API for workflow editing",
    )
    studio.add_argument(
        "--workflow",
        default=None,
        help="Path to the workflow YAML to edit (select/create in UI if not specified)",
    )
    studio.add_argument(
        "--bundle-root",
        default=None,
        help="Base directory for resolving split references",
    )
    studio.add_argument(
        "--ui-state",
        default=None,
        help="Path to the UI sidecar JSON (defaults to <workflow>.workflow-ui.json)",
    )
    studio.add_argument(
        "--workspace-root",
        default=None,
        help=(
            "Workspace root for Studio to search/create workflows. "
            "Defaults to current dir if workflow is under it, "
            "otherwise workflow parent dir, or current dir if no workflow specified."
        ),
    )
    studio.add_argument(
        "--backup-dir",
        default=".yagra/backups",
        help="Directory for storing backups",
    )
    studio.add_argument(
        "--host",
        default="127.0.0.1",
        help="Bind host",
    )
    studio.add_argument(
        "--port",
        type=int,
        default=8787,
        help="Bind port",
    )

    handlers = subparsers.add_parser(
        "handlers",
        help="Display the params schema for built-in handlers",
    )
    handlers.add_argument(
        "--format",
        choices=["text", "json"],
        default="text",
        help="Output format (default: text)",
    )

    explain = subparsers.add_parser(
        "explain",
        help="Statically analyze workflow YAML and output execution information",
    )
    explain.add_argument(
        "--workflow",
        required=True,
        help="Path to the workflow YAML to analyze",
    )
    explain.add_argument(
        "--bundle-root",
        default=None,
        help="Base directory for resolving split references",
    )
    explain.add_argument(
        "--format",
        choices=["text", "json"],
        default="json",
        help="Output format (default: json)",
    )

    mcp_parser = subparsers.add_parser(
        "mcp",
        help="Start the MCP server in stdio mode (requires yagra[mcp])",
    )
    # mcp subcommand takes no arguments
    _ = mcp_parser

    return parser


def _run_schema_command(args: argparse.Namespace) -> int:
    """Executes the `schema` subcommand.

    Outputs the JSON Schema for GraphSpec.
    Provides schema information for coding agents to generate accurate workflow YAML.

    Args:
        args: Parsed CLI arguments.

    Returns:
        Exit code. 0 on success.
    """
    schema = GraphSpec.model_json_schema()
    schema_json = json.dumps(schema, indent=2, ensure_ascii=False)

    if args.output:
        output_path = Path(args.output).expanduser().resolve()
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(schema_json + "\n", encoding="utf-8")
        print(f"schema exported: {output_path}")
    else:
        print(schema_json)
    return 0


def _run_validate_command(args: argparse.Namespace) -> int:
    """Executes the `validate` subcommand.

    Validates workflow YAML and outputs the result in text or JSON format.
    JSON format can be used for error correction loops by coding agents.
    When ``--workflow -`` is specified, reads YAML from stdin.

    Args:
        args: Parsed CLI arguments.

    Returns:
        Exit code. valid → 0, invalid → 1.
    """
    import yaml

    if args.workflow == "-":
        yaml_text = sys.stdin.read()
        try:
            payload = yaml.safe_load(yaml_text)
        except yaml.YAMLError as exc:
            issue = WorkflowValidationIssue(
                code="schema_error",
                message=f"YAML parse error: {exc}",
                location=(),
            )
            report = WorkflowValidationReport(issues=[issue])
        else:
            if not isinstance(payload, dict):
                issue = WorkflowValidationIssue(
                    code="schema_error",
                    message="workflow must be a mapping",
                    location=(),
                )
                report = WorkflowValidationReport(issues=[issue])
            else:
                report = validate_workflow_payload_for_ui(
                    payload=payload,
                    workflow_path=Path("<stdin>"),
                    bundle_root=None,
                )
    else:
        report = validate_workflow_for_ui(
            workflow_path=args.workflow,
            bundle_root=args.bundle_root,
        )

    if args.format == "json":
        report_dict = {
            "is_valid": report.is_valid,
            "issues": [issue.to_dict() for issue in report.issues],
        }
        print(json.dumps(report_dict, indent=2, ensure_ascii=False))
    else:
        print(
            format_validation_report(report), file=sys.stderr if not report.is_valid else sys.stdout
        )

    return 0 if report.is_valid else 1


def _run_visualize_command(args: argparse.Namespace) -> int:
    """Executes the `visualize` subcommand.

    Args:
        args: Parsed CLI arguments.

    Returns:
        Exit code. 0 on success.
    """
    report = validate_workflow_for_ui(
        workflow_path=args.workflow,
        bundle_root=args.bundle_root,
    )
    if not report.is_valid:
        print(format_validation_report(report), file=sys.stderr)
        return 1

    html_text = render_workflow_visualization_html(
        workflow_path=args.workflow,
        bundle_root=args.bundle_root,
        title=args.title,
    )

    output_path = Path(args.output).expanduser().resolve()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(html_text, encoding="utf-8")

    print(f"workflow visualization generated: {output_path}")
    return 0


def _run_studio_command(args: argparse.Namespace) -> int:
    """Executes the `studio` subcommand.

    Args:
        args: Parsed CLI arguments.

    Returns:
        Exit code. 0 on success.
    """
    if args.workflow is None and args.ui_state is not None:
        print(
            "--ui-state is only available when --workflow is specified.",
            file=sys.stderr,
        )
        return 2

    server = create_workflow_studio_server(
        workflow_path=args.workflow,
        bundle_root=args.bundle_root,
        ui_state_path=args.ui_state,
        workspace_root=args.workspace_root,
        backup_dir=args.backup_dir,
        host=args.host,
        port=args.port,
    )
    base_url = f"http://{args.host}:{args.port}"
    print(f"workflow studio started: {base_url}")
    print("press Ctrl+C to stop")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()
    return 0


def _run_handlers_command(args: argparse.Namespace) -> int:
    """Executes the `handlers` subcommand.

    Outputs the params schema for each built-in handler in text or JSON format.

    Args:
        args: Parsed CLI arguments.

    Returns:
        Exit code. 0 on success.
    """
    from yagra.handlers.llm_handler import LLM_HANDLER_PARAMS_SCHEMA
    from yagra.handlers.streaming_llm_handler import STREAMING_LLM_HANDLER_PARAMS_SCHEMA
    from yagra.handlers.structured_llm_handler import STRUCTURED_LLM_HANDLER_PARAMS_SCHEMA

    handlers_info = [
        {
            "name": "llm",
            "description": "LLM text output handler. Generated by create_llm_handler()",
            "params_schema": LLM_HANDLER_PARAMS_SCHEMA,
        },
        {
            "name": "structured_llm",
            "description": "Pydantic structured output handler. Generated by create_structured_llm_handler()",
            "params_schema": STRUCTURED_LLM_HANDLER_PARAMS_SCHEMA,
        },
        {
            "name": "streaming_llm",
            "description": "Streaming output handler. Generated by create_streaming_llm_handler()",
            "params_schema": STREAMING_LLM_HANDLER_PARAMS_SCHEMA,
        },
    ]

    if args.format == "json":
        print(json.dumps({"handlers": handlers_info}, indent=2, ensure_ascii=False))
    else:
        for h in handlers_info:
            print(f"## {h['name']}")
            print(f"  {h['description']}")
            print(f"  params: {json.dumps(h['params_schema'], ensure_ascii=False)}")
            print()
    return 0


def _run_explain_command(args: argparse.Namespace) -> int:
    """Executes the `explain` subcommand.

    Statically analyzes the workflow YAML and outputs execution paths,
    required handlers, and variable flow in JSON format.
    Useful for coding agents to understand a workflow before running it.

    Args:
        args: Parsed CLI arguments.

    Returns:
        Exit code. 0 on success, 1 on validation failure.
    """
    from yagra.application.use_cases.workflow_explainer import explain_workflow
    from yagra.application.use_cases.workflow_validation_reporter import (
        WorkflowValidationFailedError,
        load_validated_graph_spec,
    )

    try:
        spec = load_validated_graph_spec(
            workflow_path=args.workflow,
            bundle_root=args.bundle_root,
        )
    except WorkflowValidationFailedError as exc:
        print(format_validation_report(exc.report), file=sys.stderr)
        return 1

    result = explain_workflow(spec)

    if args.format == "json":
        print(json.dumps(result, indent=2, ensure_ascii=False))
    else:
        print(f"entry_point: {result['entry_point']}")
        print(f"exit_points: {result['exit_points']}")
        print(f"required_handlers: {result['required_handlers']}")
        print("execution_paths:")
        for path in result["execution_paths"]:
            print(f"  {' -> '.join(path)}")
        print("variable_flow:")
        for node_id, flow in result["variable_flow"].items():
            print(f"  {node_id}: inputs={flow['inputs']} outputs={flow['outputs']}")
    return 0


def _run_mcp_command(args: argparse.Namespace) -> int:
    """Executes the `mcp` subcommand.

    Starts the Yagra MCP server in stdio mode.
    Requires the 'mcp' optional dependency: uv add yagra[mcp]

    Args:
        args: Parsed CLI arguments.

    Returns:
        Exit code. 0 on success.
    """
    import asyncio

    try:
        from yagra.adapters.inbound.mcp_server import run_mcp_server
    except ImportError as exc:
        print(str(exc), file=sys.stderr)
        return 1
    asyncio.run(run_mcp_server())
    return 0


__all__ = ["Yagra", "main"]
