"""CLI application orchestration separated from the terminal UI."""

from __future__ import annotations

import argparse
import atexit
import os
import subprocess
import sys
import threading
import time
from dataclasses import dataclass

from krim.models import DEFAULT_MODELS, MissingAPIKeyError, create_model, get_api_key_env
from krim.tools import BashTool, SkillTool, create_tools, get_tool
from krim.mcp import load_mcp_config, start_mcp_servers
from krim.git import auto_commit, commit_dirty, is_git_repo, undo
from krim.ui import CliUI, RichCliUI
from krim.worktree import (
    Worktree,
    cleanup_worktrees,
    commit_worktree,
    create_worktree,
    get_current_branch,
    merge_branch,
    remove_worktree,
)
from krim_sdk.agent import Agent, AgentOptions
from krim_sdk.compaction import estimate_message_tokens
from krim_sdk.config import KrimConfig, load_config
from krim_sdk.prompt import build_system_prompt
from krim_sdk.session import SessionSnapshot, load_session
from krim_sdk.skills import discover_skills


@dataclass
class CliRuntime:
    """Resolved runtime options for a single CLI invocation."""

    provider: str
    model_name: str
    max_turns: int
    verbose: bool
    do_auto_commit: bool
    is_interactive: bool


@dataclass
class _ParallelResult:
    """Result of a single parallel task."""

    task_id: int
    prompt: str
    branch: str
    success: bool
    error: str | None = None


class SlashCommandHandler:
    """Handle interactive slash commands independently of the input loop."""

    def __init__(self, ui: CliUI):
        self.ui = ui

    def handle(self, cmd: str, agent: Agent, config: KrimConfig) -> None:
        parts = cmd.split(maxsplit=1)
        command = parts[0].lower()

        if command == "/help":
            self.ui.write(
                "[bold]Commands:[/]\n"
                "  /help      - show this help\n"
                "  /tokens    - show token usage\n"
                "  /compact   - force context compaction\n"
                "  /config    - show current config\n"
                "  /undo      - undo last krim commit\n"
                "  /verbose   - toggle verbose mode\n"
                "  exit       - quit interactive mode"
            )
            return

        if command == "/tokens":
            tokens = estimate_message_tokens(agent.messages)
            self.ui.write(f"[dim]messages: {len(agent.messages)} | ~{tokens:,} tokens[/]")
            self.ui.write(
                f"[dim]total turns: {agent.total_turns} | total tool calls: {agent.total_tool_calls}[/]"
            )
            return

        if command == "/compact":
            agent.force_compact()
            return

        if command == "/config":
            self.ui.write(f"[dim]provider: {config.provider}[/]")
            self.ui.write(f"[dim]model: {config.model or 'default'}[/]")
            self.ui.write(f"[dim]max_turns: {config.max_turns}[/]")
            self.ui.write(f"[dim]max_tokens: {config.max_tokens}[/]")
            self.ui.write(f"[dim]bash_timeout: {config.bash_timeout}s[/]")
            self.ui.write(f"[dim]token_limit: {config.token_limit:,}[/]")
            self.ui.write(f"[dim]auto_commit: {config.auto_commit}[/]")
            self.ui.write(f"[dim]safety: {'ask' if config.ask_by_default else 'off'}[/]")
            self.ui.write(f"[dim]global_dir: {config.global_dir}[/]")
            self.ui.write(f"[dim]project_dir: {config.project_dir or 'none'}[/]")
            if config.krim_md:
                self.ui.write(f"[dim]KRIM.md: {len(config.krim_md)} chars[/]")
            if config.rules:
                self.ui.write(f"[dim]rules: {len(config.rules)} loaded[/]")
            return

        if command == "/undo":
            undo()
            return

        if command == "/verbose":
            if hasattr(agent._handler, "verbose"):
                agent._handler.verbose = not agent._handler.verbose
                state = "on" if agent._handler.verbose else "off"
                self.ui.write(f"[dim]verbose: {state}[/]")
            else:
                self.ui.write("[dim]verbose toggle not supported[/]")
            return

        self.ui.write(f"[dim]unknown command: {command}. try /help[/]")


class KrimCliApp:
    """Top-level CLI controller."""

    def __init__(self, ui: CliUI | None = None):
        self.ui = ui or RichCliUI()
        self.slash_commands = SlashCommandHandler(self.ui)

    def build_runtime(
        self,
        args: argparse.Namespace,
        config: KrimConfig,
        resume_session: SessionSnapshot | None = None,
    ) -> CliRuntime:
        """Resolve CLI args into concrete runtime settings."""
        is_interactive = args.prompt is None and not args.loop and not args.parallel
        provider_was_prompted = False
        provider = args.provider or (resume_session.provider if resume_session else None)
        if provider is None:
            provider = self.ui.select_provider()
            provider_was_prompted = True

        default_model = DEFAULT_MODELS.get(provider, "gpt-5.2")
        model_name = args.model or (resume_session.model if resume_session else None)
        if model_name is None and provider_was_prompted:
            model_name = self.ui.select_model(provider, default_model)
        if model_name is None:
            model_name = default_model
        max_turns = args.max_turns if args.max_turns is not None else config.max_turns
        do_auto_commit = args.auto_commit if args.auto_commit is not None else config.auto_commit

        return CliRuntime(
            provider=provider,
            model_name=model_name,
            max_turns=max_turns,
            verbose=not args.quiet,
            do_auto_commit=do_auto_commit,
            is_interactive=is_interactive,
        )

    def resolve_prompt(self, prompt: str, skill_name: str | None, all_skills: dict) -> str:
        """Apply --skill prompt prefixing while validating the skill exists."""
        if not skill_name:
            return prompt

        if skill_name not in all_skills:
            raise ValueError(f"skill '{skill_name}' not found")

        return f"First, invoke the '{skill_name}' skill to load its instructions. Then: {prompt}"

    def run(self, args: argparse.Namespace) -> None:
        """Execute the CLI command."""
        config = load_config()

        if args.list_skills:
            self._list_skills(config)
            return

        resume_session = self._load_resume_session(args)
        if resume_session and (args.loop or args.parallel):
            self.ui.write("[red]--resume is not supported with --loop or --parallel[/]")
            sys.exit(1)

        runtime = self.build_runtime(args, config, resume_session=resume_session)

        if args.no_safety:
            config.ask_by_default = False

        self._validate_resume_provider(args, resume_session, runtime.provider)
        self._maybe_prompt_for_api_key(runtime.provider)

        self.ui.show_header(
            interactive=runtime.is_interactive,
            provider=runtime.provider,
            model_name=runtime.model_name,
            max_turns=runtime.max_turns,
            project_dir=str(config.project_dir) if config.project_dir else None,
        )
        if resume_session:
            self.ui.write(f"[dim]resumed session: {resume_session.path}[/]")
        self.ui.show_runtime(
            verbose=runtime.verbose,
            ask_by_default=config.ask_by_default,
            deny_pattern_count=len(config.deny_patterns),
            allow_command_count=len(config.allow_commands),
        )

        try:
            model = create_model(runtime.provider, runtime.model_name, max_tokens=config.max_tokens)
        except MissingAPIKeyError:
            sys.exit(1)
        all_skills = discover_skills(config.global_dir, config.project_dir)
        if all_skills:
            self.ui.write(f"[dim]skills: {len(all_skills)} available ({', '.join(all_skills.keys())})[/]")

        tools = self._create_tools(config, all_skills)
        mcp_tools, mcp_servers = self._load_mcp_tools(args, config)
        if mcp_servers:
            atexit.register(lambda: [server.stop() for server in mcp_servers])

        current_system_prompt = build_system_prompt(
            config,
            [tool.name for tool in mcp_tools],
            skills=all_skills or None,
        )
        system_prompt = (
            (resume_session.system_prompt or current_system_prompt)
            if resume_session
            else current_system_prompt
        )
        effective_log_dir = args.log_dir or (str(resume_session.path) if resume_session else None)

        if args.loop:
            if not args.prompt:
                self.ui.write("[red]--loop requires a prompt[/]")
                sys.exit(1)

            if is_git_repo():
                commit_dirty()

            try:
                prompt = self.resolve_prompt(args.prompt, args.skill, all_skills)
            except ValueError as exc:
                self.ui.write(f"[red]{exc}[/]")
                sys.exit(1)

            self._run_ralph_loop(
                prompt=prompt,
                model_factory=lambda: create_model(
                    runtime.provider, runtime.model_name, max_tokens=config.max_tokens
                ),
                system_prompt=system_prompt,
                provider=runtime.provider,
                mcp_tools=mcp_tools,
                max_turns=runtime.max_turns,
                verbose=runtime.verbose,
                max_iterations=args.max_iterations,
                exit_command=args.exit_on,
                config=config,
                all_skills=all_skills,
            )
            return

        if args.parallel:
            if not is_git_repo():
                self.ui.write("[red]--parallel requires a git repository[/]")
                sys.exit(1)

            commit_dirty()
            self._run_parallel(
                tasks=args.parallel,
                model_factory=lambda: create_model(
                    runtime.provider, runtime.model_name, max_tokens=config.max_tokens
                ),
                system_prompt=system_prompt,
                provider=runtime.provider,
                mcp_tools=mcp_tools,
                max_turns=runtime.max_turns,
                verbose=runtime.verbose,
                config=config,
                auto_merge=not args.no_merge,
                all_skills=all_skills,
            )
            return

        agent = self._create_agent(
            model=model,
            provider=runtime.provider,
            system_prompt=system_prompt,
            tools=tools,
            mcp_tools=mcp_tools,
            max_turns=runtime.max_turns,
            token_limit=config.token_limit,
            verbose=runtime.verbose,
            log_dir=effective_log_dir,
            initial_messages=resume_session.messages if resume_session else None,
        )

        if runtime.do_auto_commit and is_git_repo():
            commit_dirty()

        if args.prompt:
            try:
                prompt = self.resolve_prompt(args.prompt, args.skill, all_skills)
            except ValueError as exc:
                self.ui.write(f"[red]{exc}[/]")
                sys.exit(1)

            agent.run(prompt)
            if runtime.do_auto_commit and is_git_repo():
                auto_commit(f"krim: {args.prompt[:60]}")
            return

        self.run_interactive_session(agent, runtime, config)

    def run_interactive_session(self, agent: Agent, runtime: CliRuntime, config: KrimConfig) -> None:
        """Run the interactive prompt loop through the current UI adapter."""
        session = self.ui.create_session()
        while True:
            user_input = self.ui.prompt_input(session)
            if user_input is None:
                self.ui.show_bye(newline=True)
                break

            stripped = user_input.strip()
            if not stripped:
                continue

            if stripped.lower() in ("exit", "quit", "q"):
                self.ui.show_bye()
                break

            if stripped.startswith("/"):
                self.slash_commands.handle(stripped, agent, config)
                continue

            agent.run(user_input)

            if runtime.do_auto_commit and is_git_repo():
                auto_commit(f"krim: {stripped[:60]}")

    def _list_skills(self, config: KrimConfig) -> None:
        skills = discover_skills(config.global_dir, config.project_dir)
        if not skills:
            self.ui.write("[dim]no skills found. add them to ~/.krim/skills/ or .krim/skills/[/]")
            return

        for name, skill in skills.items():
            first_line = skill.prompt.strip().splitlines()[0] if skill.prompt.strip() else ""
            self.ui.write(f"  [bold]{name}[/]  {first_line}")

    def _create_tools(
        self,
        config: KrimConfig,
        all_skills: dict,
        *,
        cwd: str | None = None,
        ask_by_default: bool | None = None,
    ):
        tools = create_tools()
        bash_tool = get_tool(tools, "bash")
        if isinstance(bash_tool, BashTool):
            bash_tool.configure(
                deny_patterns=config.deny_patterns,
                allow_commands=config.allow_commands,
                ask_by_default=config.ask_by_default if ask_by_default is None else ask_by_default,
                max_output_chars=config.max_output_chars,
                cwd=cwd,
                default_timeout=config.bash_timeout,
            )

        skill_tool = get_tool(tools, "skill")
        if isinstance(skill_tool, SkillTool) and all_skills:
            skill_tool.configure(all_skills)

        return tools

    def _create_agent(
        self,
        *,
        model,
        provider: str,
        system_prompt: str,
        tools,
        mcp_tools: list,
        max_turns: int,
        token_limit: int,
        verbose: bool,
        log_dir: str | None = None,
        initial_messages: list[dict] | None = None,
    ) -> Agent:
        return Agent(
            model=model,
            provider=provider,
            system_prompt=system_prompt,
            tools=tools,
            mcp_tools=mcp_tools,
            options=AgentOptions(
                max_turns=max_turns,
                token_limit=token_limit,
                log_dir=log_dir,
                initial_messages=initial_messages,
            ),
            event_handler=self.ui.create_event_handler(verbose),
        )

    def _load_resume_session(self, args: argparse.Namespace) -> SessionSnapshot | None:
        if not getattr(args, "resume", None):
            return None

        try:
            return load_session(args.resume)
        except Exception as exc:
            self.ui.write(f"[red]failed to load session '{args.resume}': {exc}[/]")
            sys.exit(1)

    def _validate_resume_provider(
        self,
        args: argparse.Namespace,
        resume_session: SessionSnapshot | None,
        provider: str,
    ) -> None:
        if not resume_session or not resume_session.provider:
            return

        if args.provider and args.provider != resume_session.provider:
            self.ui.write(
                "[red]"
                f"--resume provider mismatch: session uses {resume_session.provider}, "
                f"but --provider requested {args.provider}"
                "[/]"
            )
            sys.exit(1)

        if provider != resume_session.provider:
            self.ui.write(
                "[red]"
                f"unable to infer the correct provider for resumed session; expected {resume_session.provider}"
                "[/]"
            )
            sys.exit(1)

    def _maybe_prompt_for_api_key(self, provider: str) -> None:
        env_var = get_api_key_env(provider)
        if not env_var or os.environ.get(env_var) or not sys.stdin.isatty():
            return

        label = (
            f"[yellow]{provider} requires {env_var}.[/]\n"
            "[dim]Enter it for this session only, or leave blank to cancel.[/]"
        )
        secret = self.ui.prompt_secret(label)
        if secret:
            os.environ[env_var] = secret
            self.ui.write(f"[dim]{env_var} loaded for this process only.[/]")

    def _load_mcp_tools(self, args: argparse.Namespace, config: KrimConfig) -> tuple[list, list]:
        mcp_tools = []
        mcp_servers = []
        if not args.no_mcp:
            configs = load_mcp_config(config.global_dir, config.project_dir)
            if configs:
                mcp_tools, mcp_servers = start_mcp_servers(configs)
                if mcp_tools:
                    self.ui.write(f"[dim]mcp: {len(mcp_tools)} tool(s) loaded[/]")

        return mcp_tools, mcp_servers

    def _run_exit_check(self, command: str, timeout: int = 120) -> tuple[bool, str]:
        try:
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=timeout,
            )
            output = (result.stdout + result.stderr).strip()
            return result.returncode == 0, output
        except subprocess.TimeoutExpired:
            return False, f"exit check timed out after {timeout}s"
        except Exception as exc:
            return False, f"exit check error: {exc}"

    def _run_ralph_loop(
        self,
        prompt: str,
        *,
        model_factory,
        system_prompt: str,
        provider: str,
        mcp_tools: list,
        max_turns: int,
        verbose: bool,
        max_iterations: int,
        exit_command: str | None,
        config: KrimConfig,
        all_skills: dict,
    ) -> bool:
        self.ui.write(f"\n[bold magenta]ralph loop[/] starting | max_iterations={max_iterations}")
        if exit_command:
            self.ui.write(f"[dim]exit condition: {exit_command}[/]")
        self.ui.write()

        for iteration in range(1, max_iterations + 1):
            self.ui.write(
                f"[bold magenta]━━━ ralph iteration {iteration}/{max_iterations} ━━━[/]"
            )

            tools = self._create_tools(config, all_skills)
            agent = self._create_agent(
                model=model_factory(),
                provider=provider,
                system_prompt=system_prompt,
                tools=tools,
                mcp_tools=mcp_tools,
                max_turns=max_turns,
                token_limit=config.token_limit,
                verbose=verbose,
            )
            agent.run(prompt)

            if is_git_repo():
                auto_commit(f"krim: ralph iteration {iteration}")

            if exit_command:
                self.ui.write(f"\n[dim]checking exit condition: {exit_command}[/]")
                passed, output = self._run_exit_check(exit_command)
                if output:
                    preview = output[:500]
                    if len(output) > 500:
                        preview += "..."
                    self.ui.write(f"[dim]{preview}[/]")
                if passed:
                    self.ui.write(
                        f"\n[bold green]ralph loop done![/] exit condition passed on iteration {iteration}"
                    )
                    return True

                self.ui.write("[yellow]exit condition failed, continuing...[/]\n")

            if iteration < max_iterations:
                time.sleep(2)

        self.ui.write(f"\n[bold yellow]ralph loop stopped[/] after {max_iterations} iterations")
        if exit_command:
            self.ui.write("[yellow]exit condition never passed[/]")
        return False

    def _run_agent_in_worktree(
        self,
        wt: Worktree,
        *,
        model_factory,
        system_prompt: str,
        provider: str,
        mcp_tools: list,
        max_turns: int,
        verbose: bool,
        config: KrimConfig,
        all_skills: dict,
    ) -> _ParallelResult:
        try:
            tools = self._create_tools(
                config,
                all_skills,
                cwd=str(wt.path),
                ask_by_default=False,
            )
            agent = self._create_agent(
                model=model_factory(),
                provider=provider,
                system_prompt=system_prompt,
                tools=tools,
                mcp_tools=mcp_tools,
                max_turns=max_turns,
                token_limit=config.token_limit,
                verbose=verbose,
            )
            agent.run(wt.prompt)
            commit_worktree(wt)

            return _ParallelResult(
                task_id=wt.task_id,
                prompt=wt.prompt,
                branch=wt.branch,
                success=True,
            )
        except Exception as exc:
            return _ParallelResult(
                task_id=wt.task_id,
                prompt=wt.prompt,
                branch=wt.branch,
                success=False,
                error=str(exc),
            )

    def _run_parallel(
        self,
        tasks: list[str],
        *,
        model_factory,
        system_prompt: str,
        provider: str,
        mcp_tools: list,
        max_turns: int,
        verbose: bool,
        config: KrimConfig,
        auto_merge: bool,
        all_skills: dict,
    ) -> list[_ParallelResult]:
        base_branch = get_current_branch()
        self.ui.write(f"\n[bold blue]parallel mode[/] | {len(tasks)} tasks | base: {base_branch}")

        worktrees: list[Worktree] = []
        for i, task in enumerate(tasks, 1):
            self.ui.write(f"[dim]creating worktree {i}/{len(tasks)}: {task[:60]}[/]")
            wt = create_worktree(i, task, base_branch)
            if not wt:
                self.ui.write(f"[red]failed to create worktree for task {i}[/]")
                continue
            worktrees.append(wt)

        if not worktrees:
            self.ui.write("[red]no worktrees created, aborting[/]")
            return []

        self.ui.write(f"\n[bold blue]launching {len(worktrees)} agents...[/]\n")
        results: list[_ParallelResult | None] = [None] * len(worktrees)

        def worker(idx: int, wt: Worktree) -> None:
            results[idx] = self._run_agent_in_worktree(
                wt,
                model_factory=model_factory,
                system_prompt=system_prompt,
                provider=provider,
                mcp_tools=mcp_tools,
                max_turns=max_turns,
                verbose=verbose,
                config=config,
                all_skills=all_skills,
            )

        threads = []
        for idx, wt in enumerate(worktrees):
            thread = threading.Thread(target=worker, args=(idx, wt), daemon=True)
            threads.append(thread)
            thread.start()

        for thread in threads:
            thread.join()

        self.ui.write("\n[bold blue]━━━ parallel results ━━━[/]")
        for result in results:
            if result is None:
                continue
            status = "[green]done[/]" if result.success else f"[red]failed: {result.error}[/]"
            self.ui.write(
                f"  task {result.task_id}: {status}  branch: [cyan]{result.branch}[/]"
            )

        for wt in worktrees:
            remove_worktree(wt)
        cleanup_worktrees()

        if auto_merge:
            self.ui.write(f"\n[bold blue]merging into {base_branch}...[/]")
            for result in results:
                if result is None or not result.success:
                    continue
                ok, msg = merge_branch(result.branch, into=base_branch)
                if ok:
                    self.ui.write(f"  [green]{msg}[/]")
                else:
                    self.ui.write(f"  [yellow]{msg}[/]")
        else:
            self.ui.write("\n[dim]--no-merge: branches left for manual review[/]")

        return [result for result in results if result is not None]
