from __future__ import annotations

import argparse
import os
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Literal

from tests.pipeline_setup.suite_02_pipeline.pipeline_intent_rag.helpers.env_bootstrap import PipelineEnvLoader
from tests.pipeline_setup.suite_02_pipeline.pipeline_intent_rag.helpers.models import PhraseCase, PipelineResult
from tests.pipeline_setup.suite_02_pipeline.pipeline_intent_rag.helpers.runtime import PipelineRuntime

_VALID_MODES = ("router_only", "router_rag", "full_chain")


@dataclass(slots=True, frozen=True)
class RunCliArgs:
    mode: str
    case_ids: tuple[str, ...]
    rag_session_id: str | None
    reindex_repo_path: str | None
    reindex_project_id: str | None
    test_name: str
    verbose: bool


@dataclass(slots=True, frozen=True)
class ReindexCliArgs:
    repo_path: str
    project_id: str | None


@dataclass(slots=True, frozen=True)
class ParsedCliArgs:
    command: Literal["run", "reindex"]
    payload: RunCliArgs | ReindexCliArgs


class PipelineCliRunner:
    def __init__(self, args: RunCliArgs, test_root: Path) -> None:
        self._args = args
        self._test_root = test_root

    def run(self) -> int:
        self._apply_overrides()
        runtime = PipelineRuntime(mode=self._args.mode, test_name=self._args.test_name, test_root=self._test_root)
        cases = self._filter_cases(runtime.load_cases())
        if not cases:
            print("No cases selected")
            return 2

        failed = 0
        executed = 0
        for case in cases:
            executed += 1
            try:
                result = self._run_case(runtime, case)
            except Exception as exc:
                failed += 1
                print(f"[FAIL] {case.case_id}: {exc}")
                continue
            try:
                self._validate_case(result)
                self._print_result(result, status="OK", reason=None, force_diagnostics=self._args.verbose)
            except Exception as exc:
                failed += 1
                self._print_result(result, status="FAIL", reason=str(exc), force_diagnostics=True)

        print("-")
        print(f"Mode: {self._args.mode}")
        print(f"Executed: {executed}")
        print(f"Failed: {failed}")
        print(f"Artifact: {runtime.artifact_path}")
        return 0 if failed == 0 else 1

    def _apply_overrides(self) -> None:
        if self._args.rag_session_id:
            os.environ["INTENT_PIPELINE_FORCE_RAG_SESSION_ID"] = self._args.rag_session_id
        if self._args.reindex_repo_path:
            os.environ["INTENT_PIPELINE_REINDEX_REPO_PATH"] = self._args.reindex_repo_path
        if self._args.reindex_project_id:
            os.environ["INTENT_PIPELINE_REINDEX_PROJECT_ID"] = self._args.reindex_project_id

    def _filter_cases(self, cases: list[PhraseCase]) -> list[PhraseCase]:
        if not self._args.case_ids:
            return cases
        wanted = set(self._args.case_ids)
        return [case for case in cases if case.case_id in wanted]

    def _run_case(self, runtime: PipelineRuntime, case: PhraseCase) -> PipelineResult:
        if self._args.mode == "router_only":
            return runtime.run_router_only(case)
        if self._args.mode == "router_rag":
            return runtime.run_router_rag(case)
        return runtime.run_full_chain(case)

    def _validate_case(self, result: PipelineResult) -> None:
        case = result.case
        if case.expected_intent and result.intent != case.expected_intent:
            raise AssertionError(f"expected_intent={case.expected_intent} actual={result.intent}")
        if self._args.mode in {"router_rag", "full_chain"} and case.expect_non_empty_rag and not result.rag_rows:
            raise AssertionError(f"RAG returned empty list for rag_session_id={result.rag_session_id}")
        if self._args.mode == "full_chain" and not (result.llm_answer or "").strip():
            raise AssertionError("LLM answer must not be empty")

    def _print_result(
        self,
        result: PipelineResult,
        *,
        status: str,
        reason: str | None,
        force_diagnostics: bool,
    ) -> None:
        record = result.to_record()
        if status != "OK":
            print(f"status: {status}")
        print(f"case_id: {record['case_id']}")
        print(f"expected_intent: {record['expected_intent']}")
        print(f"actual_intent: {record['actual_intent']}")
        print("summary:")
        self._print_block(record.get("summary"), indent=2)
        if force_diagnostics:
            if reason:
                print(f"failure_reason: {reason}")
            print("diagnostics:")
            self._print_block(record.get("diagnostics"), indent=2)
        print("-")

    def _print_block(self, value, *, indent: int) -> None:
        prefix = " " * indent
        if isinstance(value, dict):
            for key, item in value.items():
                if isinstance(item, (dict, list)):
                    print(f"{prefix}{key}:")
                    self._print_block(item, indent=indent + 2)
                else:
                    print(f"{prefix}{key}: {item}")
            return
        if isinstance(value, list):
            for item in value:
                if isinstance(item, dict):
                    print(f"{prefix}-")
                    self._print_block(item, indent=indent + 2)
                else:
                    print(f"{prefix}- {item}")
            return
        print(f"{prefix}{value}")


class ReindexCliRunner:
    def __init__(self, args: ReindexCliArgs, test_root: Path) -> None:
        self._args = args
        self._test_root = test_root

    def run(self) -> int:
        try:
            PipelineEnvLoader(self._test_root).load()
            repo_path = Path(self._args.repo_path).expanduser().resolve()
            if not repo_path.is_dir():
                print(f"[FAIL] repository path is not a directory: {repo_path}")
                return 2
            from app.modules.rag.persistence.repository import RagRepository
            from tests.pipeline_setup.utils.rag_indexer import RagSessionIndexer

            RagSessionIndexer(RagRepository()).index_repo(repo_path, project_id=self._args.project_id)
            return 0
        except Exception as exc:
            print(f"[FAIL] reindex error: {exc}")
            return 1


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Run intent-router pipeline CLI")
    sub = parser.add_subparsers(dest="command")

    run_cmd = sub.add_parser("run", help="Run test pipeline")
    run_cmd.add_argument("--mode", choices=_VALID_MODES, required=True, help="Execution mode")
    run_cmd.add_argument("--case-id", action="append", default=[], help="Case id to run (repeatable)")
    run_cmd.add_argument("--rag-session-id", default=None, help="Force rag_session_id for all cases")
    run_cmd.add_argument("--reindex-repo-path", default=None, help="Path to local repository for reindex before retrieval")
    run_cmd.add_argument("--reindex-project-id", default=None, help="project_id for newly indexed rag_session")
    run_cmd.add_argument("--test-name", default=None, help="Artifact prefix (defaults to cli_<mode>)")
    run_cmd.add_argument(
        "--verbose",
        "--debug",
        dest="verbose",
        action="store_true",
        help="Print diagnostics for every case",
    )

    reindex_cmd = sub.add_parser("reindex", help="Create a new RAG session by indexing local repository")
    reindex_cmd.add_argument("--repo-path", required=True, help="Path to local repository directory")
    reindex_cmd.add_argument("--project-id", default=None, help="Project id for created rag_session")

    return parser


def _parse_args(argv: list[str] | None = None) -> ParsedCliArgs:
    parser = _build_parser()
    raw = list(sys.argv[1:] if argv is None else argv)
    if raw and raw[0] not in {"run", "reindex"} and raw[0] not in {"-h", "--help"}:
        raw = ["run", *raw]
    ns = parser.parse_args(raw)
    command = str(getattr(ns, "command", "") or "")
    if command == "run":
        mode = str(ns.mode)
        test_name = (ns.test_name or f"cli_{mode}").strip() or f"cli_{mode}"
        payload = RunCliArgs(
            mode=mode,
            case_ids=tuple(str(item).strip() for item in ns.case_id if str(item).strip()),
            rag_session_id=str(ns.rag_session_id).strip() or None if ns.rag_session_id is not None else None,
            reindex_repo_path=str(ns.reindex_repo_path).strip() or None if ns.reindex_repo_path is not None else None,
            reindex_project_id=str(ns.reindex_project_id).strip() or None if ns.reindex_project_id is not None else None,
            test_name=test_name,
            verbose=bool(getattr(ns, "verbose", False)),
        )
        return ParsedCliArgs(command="run", payload=payload)
    if command == "reindex":
        payload = ReindexCliArgs(
            repo_path=str(ns.repo_path).strip(),
            project_id=str(ns.project_id).strip() or None if ns.project_id is not None else None,
        )
        return ParsedCliArgs(command="reindex", payload=payload)
    parser.error("command is required: run or reindex")


def main(argv: list[str] | None = None) -> int:
    try:
        args = _parse_args(argv)
        test_root = Path(__file__).resolve().parent
        if args.command == "run":
            return PipelineCliRunner(args=args.payload, test_root=test_root).run()
        return ReindexCliRunner(args=args.payload, test_root=test_root).run()
    except KeyboardInterrupt:
        print("Interrupted")
        return 130


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))
