"""
CQ-AI Interactive REPL
Implements a Claude Code-style interactive command loop.
"""

import os
import sys
import shlex
from pathlib import Path
from typing import Optional

try:
    from prompt_toolkit import PromptSession
    from prompt_toolkit.history import InMemoryHistory
    HAS_PROMPT_TOOLKIT = True
except ImportError:
    HAS_PROMPT_TOOLKIT = False

# Fix Windows console encoding
if sys.platform == "win32":
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")

from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.markdown import Markdown
from rich.prompt import Prompt, Confirm
import yaml
from dotenv import set_key, load_dotenv

# Internal imports - assuming package structure remains ai_coder for now, 
# even if package name changes to cq-ai on PyPI
from ai_coder.main import load_config, get_api_key
from ai_coder.llm.interface import LLMConfig, get_provider
from ai_coder.agent.orchestrator import AgentOrchestrator, AgentConfig
from ai_coder.tools.base import create_tool_registry
from ai_coder.agents.launcher import AgentLauncher
from ai_coder.agents.base import AgentType
from ai_coder import __version__

console = Console()

class InteractiveREPL:
    def __init__(self, project_dir: Optional[Path] = None, provider: str = None, model: str = None, api_key: str = None):
        self.project_dir = (project_dir or Path.cwd()).resolve()
        self.config = load_config()
        
        # LLM Setup
        llm_config = self.config.get("llm", {})
        self.provider = provider or llm_config.get("provider", "openai")
        self.model = model or llm_config.get("model", "gpt-4o")
        self.api_key = api_key
        
        # Session state
        self.tool_registry = create_tool_registry(self.project_dir)
        self.llm = None
        self._init_llm()

    def _init_llm(self):
        """Initialize LLM provider."""
        try:
            # Resolve API Key
            keyless_providers = {"ollama", "lmstudio", "custom"}
            
            # If provider explicitly passed, use it. Otherwise config. 
            # If config missing or key missing for selected provider, run onboarding.
            
            if not self.api_key and self.provider not in keyless_providers:
                try:
                    # Check if we have a key in env
                    get_api_key(self.provider)
                except (SystemExit, Exception):
                    # Key missing -> Run Onboarding
                    self._run_onboarding()
            
            # Reload config after onboarding
            self.config = load_config()
            llm_config = self.config.get("llm", {})
            self.provider = self.provider or llm_config.get("provider", "openai")
            self.model = self.model or llm_config.get("model", "gpt-4o")
            
            # Get key again
            if self.provider in keyless_providers:
                key = self.api_key or ""
            else:
                try:
                    key = self.api_key or get_api_key(self.provider)
                except (SystemExit, Exception):
                     # Should not happen if onboarding worked, but just in case
                     console.print(f"[yellow]API key for {self.provider} still not found.[/yellow]")
                     key = console.input(f"[bold]Enter {self.provider} API Key: [/bold]").strip()

            api_base = self.config.get("llm", {}).get("api_base") or os.getenv(f"{self.provider.upper()}_API_BASE")

            llm_cfg = LLMConfig(
                model=self.model,
                max_tokens=self.config.get("llm", {}).get("max_tokens", 4096),
                temperature=0.1,
                api_key=key,
                api_base=api_base,
            )
            self.llm = get_provider(self.provider, llm_cfg)
            
        except Exception as e:
            console.print(f"[red]Error initializing LLM: {e}[/red]")
            console.print("[yellow]Continuing with limited functionality...[/yellow]")

    def _run_onboarding(self):
        """Run interactive setup wizard."""
        console.clear()
        console.print(Panel("[bold green]Welcome to CQ-AI Setup[/bold green]", title="Initial Configuration"))
        console.print("Let's set up your AI provider.\n")
        
        providers = ["openai", "anthropic", "huggingface", "ollama", "google", "groq", "custom"]
        console.print("[bold]Select LLM Provider:[/bold]")
        for i, p in enumerate(providers, 1):
            console.print(f"  {i}. {p}")
            
        choice = Prompt.ask("Choice", choices=[str(i) for i in range(1, len(providers)+1)], default="1")
        selected_provider = providers[int(choice)-1]
        
        key = ""
        keyless = {"ollama", "lmstudio"}
        if selected_provider not in keyless:
            key = Prompt.ask(f"Enter {selected_provider} API Key", password=True)
        
        # Default models
        default_models = {
            "openai": "gpt-4o",
            "anthropic": "claude-3-5-sonnet-20240620",
            "huggingface": "Qwen/Qwen2.5-Coder-32B-Instruct", 
            "ollama": "llama3",
            "google": "gemini-1.5-pro",
            "groq": "llama3-70b-8192",
        }
        default_model = default_models.get(selected_provider, "gpt-4o")
        model = Prompt.ask("Model", default=default_model)
        
        # Save
        if Confirm.ask("Save configuration?"):
            self._save_config(selected_provider, model)
            if key:
                self._save_env(selected_provider, key)
            
            # Update current session
            self.provider = selected_provider
            self.model = model
            self.api_key = key
            console.print("[green]Configuration saved![/green]\n")

    def _save_config(self, provider, model):
        """Save config to .ai-coder/config.yaml"""
        config_dir = Path.home() / ".ai-coder"
        config_dir.mkdir(exist_ok=True)
        config_path = config_dir / "config.yaml"
        
        if config_path.exists():
            with open(config_path, "r") as f:
                data = yaml.safe_load(f) or {}
        else:
            data = {}
            
        if "llm" not in data: data["llm"] = {}
        data["llm"]["provider"] = provider
        data["llm"]["model"] = model
        
        with open(config_path, "w") as f:
            yaml.dump(data, f)
            
    def _save_env(self, provider, key):
        """Save API key to .env file in project or home."""
        # Prefer local .env if exists, else home .env? 
        # Actually easier to just valid local .env for now or export.
        # But set_key writes to a file.
        
        env_path = self.project_dir / ".env"
        # If it doesn't exist, maybe create it?
        if not env_path.exists():
            env_path.touch()
            
        # Map provider to env var
        var_name = f"{provider.upper()}_API_KEY"
        set_key(env_path, var_name, key)
        os.environ[var_name] = key # Set for current process too

    def run_loop(self):
        """Main REPL loop."""
        self._print_banner()
        
        # Initialize prompt session if available
        session = None
        if HAS_PROMPT_TOOLKIT:
            session = PromptSession(history=InMemoryHistory())
        
        while True:
            try:
                # Dynamic prompt showing generic info or context
                cwd_name = self.project_dir.name
                prompt_text = f"cq-ai ({cwd_name})> "
                
                if session:
                    # Use prompt_toolkit for history support
                    # We need to print formatted prompt part first if we want colors, 
                    # but prompt_toolkit handles prompt text differently. 
                    # For simplicity, let's use plain text prompt with prompt_toolkit 
                    # or use HTML for coloring if meaningful.
                    from prompt_toolkit.formatted_text import HTML
                    user_input = session.prompt(HTML(f"<ansigreen>cq-ai</ansigreen> <ansiblue>({cwd_name})</ansiblue>&gt; "))
                else:
                    user_input = console.input(f"[bold green]cq-ai[/bold green] [blue]({cwd_name})[/blue]> ")
                
                if not user_input.strip():
                    continue
                
                if user_input.startswith("/"):
                    self._handle_command(user_input)
                else:
                    self._handle_natural_language(user_input)
                    
            except KeyboardInterrupt:
                # Handle Ctrl+C gracefully - assume user wants to cancel current input line
                continue
            except EOFError:
                # Ctrl+D to exit
                break
            except Exception as e:
                console.print(f"[red]Error: {e}[/red]")

    def _print_banner(self):
        console.clear()
        console.print(Panel(
            f"[bold green]CQ-AI v{__version__}[/bold green]\n"
            f"[dim]Developed by CQ Team[/dim]\n"
            f"Provider: {self.provider} | Model: {self.model}\n"
            f"Project: {self.project_dir}\n\n"
            "Type your request in plain English, or use slash commands.\n"
            "Type [bold]/help[/bold] for list of commands.",
            border_style="green",
            title="Welcome",
            title_align="left"
        ))

    def _handle_command(self, user_input: str):
        parts = shlex.split(user_input)
        cmd = parts[0].lower()
        args = " ".join(parts[1:]) if len(parts) > 1 else ""
        
        commands = {
            "/exit": self._exit,
            "/quit": self._exit,
            "/help": self._help,
            "/clear": self._clear,
            "/build": lambda: self._run_agent_command("build", args),
            "/debug": lambda: self._run_agent_command("debug", args),
            "/review": lambda: self._run_agent_command("review", args),
            "/plan": lambda: self._run_agent_command("plan", args),
            "/refactor": lambda: self._run_agent_command("refactor", args),
            "/security": lambda: self._run_agent_command("security", args),
            "/test": lambda: self._run_agent_command("tdd", args),
            "/fix": lambda: self._run_agent_command("build-fix", args),
            "/config": self._run_onboarding,
        }
        
        if cmd in commands:
            commands[cmd]()
        else:
            console.print(f"[red]Unknown command: {cmd}[/red]")

    def _handle_natural_language(self, prompt: str):
        """Send natural language request to Orchestrator."""
        if not self.llm:
            console.print("[red]LLM not initialized.[/red]")
            return

        console.print(f"[dim]Thinking...[/dim]")
        
        # Use AgentOrchestrator for general tasks as it's the most capable
        agent_cfg = AgentConfig(
            max_iterations=30,
            auto_confirm=True, # Interactive mode implies user is watching? Maybe toggle.
            provider=self.provider,
            model=self.model,
            verbose=False
        )
        
        # For now, default to auto_confirm=False for safety unless configured otherwise
        # But wait, ai-coder default was False. 
        # Claude Code is usually "Human in the loop" but fast.
        # Let's stick to config defaults.
        
        orchestrator = AgentOrchestrator(
            llm=self.llm,
            project_root=self.project_dir,
            config=agent_cfg
        )
        
        try:
            orchestrator.run(prompt)
        except Exception as e:
            console.print(f"[red]Agent error: {e}[/red]")

    def _run_agent_command(self, mode: str, args: str):
        """Run a specific agent workflow."""
        if not self.llm:
            console.print("[red]LLM not initialized.[/red]")
            return

        if not args:
            # Prompt user for missing args if needed
            if mode == "build":
                 args = console.input("[bold]What do you want to build? [/bold]")
            elif mode == "debug":
                 args = console.input("[bold]What is the issue? [/bold]")
            # ... others might default to "current context"

        if not args and mode not in ["review", "security", "fix"]:
            console.print("[yellow]Cancelled (no input).[/yellow]")
            return

        # Map mode to agent types or logic
        # This logic mirrors main.py commands somewhat
        
        try:
            if mode == "build":
                # Build uses Orchestrator with special prompt
                task = f"Build project: {args}"
                # logic similar to main.py build command...
                self._handle_natural_language(task)
                return

            launcher = AgentLauncher(self.llm, self.tool_registry, self.project_dir)
            
            agent_type_map = {
                "debug": AgentType.DEBUGGER,
                "review": AgentType.CODE_REVIEWER,
                "plan": AgentType.PLANNER,
                "refactor": AgentType.REFACTOR,
                "security": AgentType.SECURITY_REVIEWER,
                "tdd": AgentType.TDD_GUIDE,
                "build-fix": AgentType.BUILD_FIXER
            }
            
            target_agent = agent_type_map.get(mode)
            if target_agent:
                result = launcher.launch_single(target_agent, args or "Analyze current codebase")
                if result.success:
                     console.print(result.output)
                else:
                     console.print(f"[red]{result.error}[/red]")
            else:
                console.print(f"[red]Method for {mode} not implemented yet.[/red]")

        except Exception as e:
            console.print(f"[red]Execution error: {e}[/red]")

    def _exit(self):
        console.print("[green]Goodbye![/green]")
        sys.exit(0)

    def _clear(self):
        os.system('cls' if os.name == 'nt' else 'clear')
        self._print_banner()

    def _help(self):
        md = """
# CQ-AI Commands

| Command | Description |
|---------|-------------|
| `/build [desc]` | Build a new project or feature |
| `/debug [issue]` | Debug an issue |
| `/review [path]` | Review code quality |
| `/refactor [path]` | Refactor code |
| `/plan [task]` | Create implementation plan |
| `/security [path]`| Run security audit |
| `/test [feature]` | TDD implementation |
| `/fix [cmd]` | Fix build/run errors |
| `/config` | Run setup wizard |
| `/clear` | Clear screen |
| `/quit` | Exit |

**Just type natural language** to use the general assistant.
        """
        console.print(Markdown(md))

def launch_interactive(project_dir=None, provider=None, model=None, api_key=None):
    repl = InteractiveREPL(project_dir, provider, model, api_key)
    repl.run_loop()
