"""krim - thin CLI agent. trust the model, keep the harness light.

Usage:
  krim "fix the bug in main.py"
  krim --provider claude --model claude-sonnet-4-5-20250929 "refactor this"
  krim --provider openai --model gpt-5.2 "add tests"
  krim --max-turns 20 "big refactor task"
  krim --skill deploy "ship it"
  krim --auto-commit "fix and commit"
  krim --loop --exit-on "pytest" "implement the PRD"
  krim --parallel "add auth" "add payments" "add tests"
  krim                         # interactive mode
"""

from __future__ import annotations

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

from rich.console import Console

from krim import __version__
from krim.models import create_model, DEFAULT_MODELS, MissingAPIKeyError
from krim.tools import create_tools, get_tool, SkillTool
from krim.tools import BashTool
from krim.mcp import load_mcp_config, start_mcp_servers
from krim.git import is_git_repo, commit_dirty, auto_commit, undo
from krim.worktree import (
    create_worktree, commit_worktree, remove_worktree,
    cleanup_worktrees, merge_branch, get_current_branch, Worktree,
)
from krim.ui import print_banner, print_banner_oneliner, create_session, prompt_input, select_provider

# Import from SDK
from krim_sdk.config import load_config, KrimConfig
from krim_sdk.agent import Agent
from krim_sdk.skills import discover_skills
from krim_sdk.prompt import build_system_prompt
from krim_sdk.compaction import estimate_message_tokens

console = Console()


def _handle_slash_command(cmd: str, agent: Agent, config: KrimConfig, verbose: bool):
    """Handle interactive slash commands."""
    parts = cmd.split(maxsplit=1)
    command = parts[0].lower()

    if command == "/help":
        console.print(
            "[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"
        )
    elif command == "/tokens":
        tokens = estimate_message_tokens(agent.messages)
        console.print(f"[dim]messages: {len(agent.messages)} | ~{tokens:,} tokens[/]")
        console.print(f"[dim]total turns: {agent.total_turns} | total tool calls: {agent.total_tool_calls}[/]")
    elif command == "/compact":
        agent.force_compact()
    elif command == "/config":
        console.print(f"[dim]provider: {config.provider}[/]")
        console.print(f"[dim]model: {config.model or 'default'}[/]")
        console.print(f"[dim]max_turns: {config.max_turns}[/]")
        console.print(f"[dim]max_tokens: {config.max_tokens}[/]")
        console.print(f"[dim]bash_timeout: {config.bash_timeout}s[/]")
        console.print(f"[dim]token_limit: {config.token_limit:,}[/]")
        console.print(f"[dim]auto_commit: {config.auto_commit}[/]")
        console.print(f"[dim]safety: {'ask' if config.ask_by_default else 'off'}[/]")
        console.print(f"[dim]global_dir: {config.global_dir}[/]")
        console.print(f"[dim]project_dir: {config.project_dir or 'none'}[/]")
        if config.krim_md:
            console.print(f"[dim]KRIM.md: {len(config.krim_md)} chars[/]")
        if config.rules:
            console.print(f"[dim]rules: {len(config.rules)} loaded[/]")
    elif command == "/undo":
        undo()
    elif command == "/verbose":
        agent.verbose = not agent.verbose
        console.print(f"[dim]verbose: {'on' if agent.verbose else 'off'}[/]")
    else:
        console.print(f"[dim]unknown command: {command}. try /help[/]")


def parse_args():
    p = argparse.ArgumentParser(
        prog="krim",
        description="Thin CLI agent. Trust the model, keep the harness light.",
    )
    p.add_argument("prompt", nargs="?", help="task to perform (omit for interactive mode)")
    p.add_argument("--provider", "-p", default=None, choices=["claude", "openai", "vertex_ai"],
                   help="model provider (default: from config or claude)")
    p.add_argument("--model", "-m", default=None,
                   help="model name (default: provider's default)")
    p.add_argument("--max-turns", "-t", type=int, default=None,
                   help="max agent turns (default: from config or 10)")
    p.add_argument("--skill", "-s", default=None,
                   help="activate a skill by name")
    p.add_argument("--list-skills", action="store_true",
                   help="list available skills")
    p.add_argument("--no-mcp", action="store_true",
                   help="disable MCP servers")
    p.add_argument("--auto-commit", action="store_true", default=None,
                   help="auto-commit after agent edits")
    p.add_argument("--no-safety", action="store_true",
                   help="disable bash safety prompts (auto-allow all)")
    p.add_argument("--quiet", "-q", action="store_true",
                   help="hide turn logs, show only final response")
    p.add_argument("--log-dir", default=None,
                   help="save conversation logs to directory (for benchmarks)")
    p.add_argument("--version", "-v", action="version", version=f"krim {__version__}")

    # ralph loop options
    p.add_argument("--loop", action="store_true",
                   help="run in loop mode (ralph wiggum loop)")
    p.add_argument("--max-iterations", type=int, default=20,
                   help="max loop iterations before stopping (default: 20)")
    p.add_argument("--exit-on", default=None,
                   help="shell command that must exit 0 to stop the loop (e.g. 'pytest')")

    # parallel worktree options
    p.add_argument("--parallel", nargs="+", metavar="TASK",
                   help="run multiple tasks in parallel via git worktrees")
    p.add_argument("--no-merge", action="store_true",
                   help="keep parallel branches separate (don't auto-merge)")

    return p.parse_args()


# -- ralph wiggum loop --

def _run_exit_check(command: str, timeout: int = 120) -> tuple[bool, str]:
    """Run the exit condition command. Returns (passed, output)."""
    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 e:
        return False, f"exit check error: {e}"


def _run_ralph_loop(
    prompt: str,
    *,
    model_factory,
    system_prompt: str,
    provider: str,
    tools_factory,
    mcp_tools: list,
    max_turns: int,
    verbose: bool,
    max_iterations: int,
    exit_command: str | None,
    config: KrimConfig,
    all_skills: dict,
):
    """Run the ralph wiggum loop: fresh agent per iteration, git as memory."""
    console.print(f"\n[bold magenta]ralph loop[/] starting | max_iterations={max_iterations}")
    if exit_command:
        console.print(f"[dim]exit condition: {exit_command}[/]")
    console.print()

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

        # fresh tools + agent each iteration (clean context)
        tools = tools_factory()
        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,
                max_output_chars=config.max_output_chars,
                default_timeout=config.bash_timeout,
            )

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

        agent = Agent(
            model=model_factory(),
            provider=provider,
            system_prompt=system_prompt,
            tools=tools,
            mcp_tools=mcp_tools,
            max_turns=max_turns,
            verbose=verbose,
            token_limit=config.token_limit,
        )

        # run agent
        agent.run(prompt)

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

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

            console.print(f"[yellow]exit condition failed, continuing...[/]\n")

        # small delay between iterations
        if iteration < max_iterations:
            time.sleep(2)

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


# -- parallel worktree execution --

@dataclass
class _ParallelResult:
    """Result of a single parallel task."""
    task_id: int
    prompt: str
    branch: str
    success: bool
    error: str | None = None


def _run_agent_in_worktree(
    wt: Worktree,
    *,
    model_factory,
    system_prompt: str,
    provider: str,
    mcp_tools: list,
    max_turns: int,
    verbose: bool,
    config: KrimConfig,
    all_skills: dict,
) -> _ParallelResult:
    """Run a single agent inside a worktree directory. Called from a thread."""
    try:
        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=False,  # non-interactive in parallel
                max_output_chars=config.max_output_chars,
                cwd=str(wt.path),
                default_timeout=config.bash_timeout,
            )

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

        agent = Agent(
            model=model_factory(),
            provider=provider,
            system_prompt=system_prompt,
            tools=tools,
            mcp_tools=mcp_tools,
            max_turns=max_turns,
            verbose=verbose,
            token_limit=config.token_limit,
        )

        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 e:
        return _ParallelResult(
            task_id=wt.task_id,
            prompt=wt.prompt,
            branch=wt.branch,
            success=False,
            error=str(e),
        )


def _run_parallel(
    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]:
    """Run multiple tasks in parallel using git worktrees."""
    base_branch = get_current_branch()
    console.print(f"\n[bold blue]parallel mode[/] | {len(tasks)} tasks | base: {base_branch}")

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

    if not worktrees:
        console.print("[red]no worktrees created, aborting[/]")
        return []

    console.print(f"\n[bold blue]launching {len(worktrees)} agents...[/]\n")

    # run agents in threads
    results: list[_ParallelResult] = [None] * len(worktrees)  # type: ignore

    def _worker(idx: int, wt: Worktree):
        results[idx] = _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):
        t = threading.Thread(target=_worker, args=(idx, wt), daemon=True)
        threads.append(t)
        t.start()

    # wait for all threads
    for t in threads:
        t.join()

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

    # cleanup worktrees (keep branches)
    for wt in worktrees:
        remove_worktree(wt)
    cleanup_worktrees()

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

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


def main():
    args = parse_args()

    # load layered config
    config = load_config()

    # determine if interactive mode
    is_interactive = args.prompt is None and not args.loop and not args.parallel

    # CLI flags override config
    # Without --provider, always show selection menu
    if args.provider is None:
        provider = select_provider()
    else:
        provider = args.provider
    # Use --model if specified, otherwise use provider's default (not config.model)
    model_name = args.model or DEFAULT_MODELS.get(provider, "gpt-5.2")
    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

    verbose = not args.quiet  # verbose by default, --quiet to hide turn logs
    if args.no_safety:
        config.ask_by_default = False

    # list skills
    if args.list_skills:
        skills = discover_skills(config.global_dir, config.project_dir)
        if not skills:
            console.print("[dim]no skills found. add them to ~/.krim/skills/ or .krim/skills/[/]")
        for name, skill in skills.items():
            first_line = skill.prompt.strip().splitlines()[0] if skill.prompt.strip() else ""
            console.print(f"  [bold]{name}[/]  {first_line}")
        return

    # header - full banner for interactive, one-liner for single prompt
    if is_interactive:
        print_banner(provider, model_name, max_turns, str(config.project_dir) if config.project_dir else None)
    else:
        print_banner_oneliner(provider, model_name, max_turns)
        if config.project_dir:
            console.print(f"[dim]config: {config.project_dir}[/]")
    if verbose:
        console.print(f"[dim]verbose: on | safety: {'ask' if config.ask_by_default else 'off'}[/]")
        console.print(f"[dim]deny_patterns: {len(config.deny_patterns)} | allow_commands: {len(config.allow_commands)}[/]")

    # create model
    model = create_model(provider, model_name, max_tokens=config.max_tokens)

    # discover and configure skills
    all_skills = discover_skills(config.global_dir, config.project_dir)
    if all_skills:
        console.print(f"[dim]skills: {len(all_skills)} available ({', '.join(all_skills.keys())})[/]")

    # create tools and configure bash safety
    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,
            max_output_chars=config.max_output_chars,
            default_timeout=config.bash_timeout,
        )

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

    # load MCP tools
    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:
                console.print(f"[dim]mcp: {len(mcp_tools)} tool(s) loaded[/]")

    # cleanup MCP servers on exit
    if mcp_servers:
        atexit.register(lambda: [s.stop() for s in mcp_servers])

    # build system prompt with skill metadata (progressive disclosure)
    extra_tool_names = [t.name for t in mcp_tools]
    system_prompt = build_system_prompt(config, extra_tool_names, skills=all_skills or None)

    # -- ralph wiggum loop mode --
    if args.loop:
        if not args.prompt:
            console.print("[red]--loop requires a prompt[/]")
            sys.exit(1)

        # git: protect uncommitted changes
        if is_git_repo():
            commit_dirty()

        prompt = args.prompt
        if args.skill:
            if args.skill not in all_skills:
                console.print(f"[red]skill '{args.skill}' not found[/]")
                sys.exit(1)
            prompt = f"First, invoke the '{args.skill}' skill to load its instructions. Then: {prompt}"

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

    # -- parallel worktree mode --
    if args.parallel:
        if not is_git_repo():
            console.print("[red]--parallel requires a git repository[/]")
            sys.exit(1)

        # protect uncommitted changes
        commit_dirty()

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

    # create agent
    agent = Agent(
        model=model,
        provider=provider,
        system_prompt=system_prompt,
        tools=tools,
        mcp_tools=mcp_tools,
        max_turns=max_turns,
        verbose=verbose,
        token_limit=config.token_limit,
        log_dir=args.log_dir,
    )

    # git: protect uncommitted changes
    if do_auto_commit and is_git_repo():
        commit_dirty()

    # single prompt mode
    if args.prompt:
        prompt = args.prompt
        # --skill flag: prepend skill invocation instruction
        if args.skill:
            if args.skill not in all_skills:
                console.print(f"[red]skill '{args.skill}' not found[/]")
                sys.exit(1)
            prompt = f"First, invoke the '{args.skill}' skill to load its instructions. Then: {prompt}"
        agent.run(prompt)
        if do_auto_commit and is_git_repo():
            auto_commit(f"krim: {args.prompt[:60]}")
        return

    # interactive mode
    session = create_session()
    while True:
        user_input = prompt_input(session)
        if user_input is None:
            console.print("\n[dim]bye[/]")
            break

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

        if stripped.lower() in ("exit", "quit", "q"):
            console.print("[dim]bye[/]")
            break

        # slash commands
        if stripped.startswith("/"):
            _handle_slash_command(stripped, agent, config, verbose)
            continue

        agent.run(user_input)

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


if __name__ == "__main__":
    main()
