#!/usr/bin/env python3
"""nanocode-dspy - minimal claude code alternative using DSPy ReAct"""

import os
from modaic import PrecompiledProgram, PrecompiledConfig
import dspy
from dspy.utils.callback import BaseCallback

from utils import (
    RESET, BOLD, DIM, BLUE, CYAN, GREEN, RED, MAGENTA,
    separator, render_markdown,
    read_file, write_file, edit_file, glob_files, grep_files, run_bash,
    AVAILABLE_MODELS, select_model,
)


# --- DSPy Signature ---

class CodingAssistant(dspy.Signature):
    """You are a concise coding assistant. Help the user with their coding task by using the available tools to read, write, edit files, search the codebase, and run commands."""

    task: str = dspy.InputField(desc="The user's coding task or question")
    answer: str = dspy.OutputField(desc="Your response to the user after completing the task")
    affected_files: list[str] = dspy.OutputField(desc="List of files that were written or modified during the task")

# ReAct agent with tools

tools = [read_file, write_file, edit_file, glob_files, grep_files, run_bash]


class ToolLoggingCallback(BaseCallback):
    """Callback that logs tool calls as they happen."""
    
    def on_tool_start(self, call_id, instance, inputs):
        """Log when a tool starts executing."""
        tool_name = instance.name if hasattr(instance, 'name') else str(instance)
        # Format args nicely
        args_str = ", ".join(f"{k}={repr(v)[:50]}" for k, v in inputs.items())
        print(f"  {MAGENTA}⏺ {tool_name}({args_str}){RESET}", flush=True)
    
    def on_tool_end(self, call_id, outputs, exception):
        """Log when a tool finishes executing."""
        if exception:
            print(f"    {RED}Error: {exception}{RESET}", flush=True)
    
    def on_module_end(self, call_id, outputs, exception):
        """Log when the finish tool is called (ReAct completion)."""
        # Check if this is a ReAct prediction with tool_calls
        if outputs and 'tool_calls' in outputs:
            for call in outputs['tool_calls']:
                args_str = ", ".join(f"{k}={repr(v)[:50]}" for k, v in call.args.items())
                if call.name == 'finish':
                    print(f"  {GREEN}⏺ finish{RESET}", flush=True)
                else:
                    print(f"  {MAGENTA}⏺ {call.name}({args_str}){RESET}", flush=True)


class AgentConfig(PrecompiledConfig):
    max_iters: int = 15
    lm: str = "openrouter/anthropic/claude-3.5-sonnet"  # Default fallback
    api_base: str = "https://openrouter.ai/api/v1"
    max_tokens: int = 8192

class AgentProgram(PrecompiledProgram):
    config: AgentConfig
    
    def __init__(self, config: AgentConfig, **kwargs):
        self.config = config
        super().__init__(config, **kwargs)

        # Configure logging callback globally
        dspy.settings.configure(callbacks=[ToolLoggingCallback()])
        
        agent = dspy.ReAct(CodingAssistant, tools=tools, max_iters=self.config.max_iters)
        lm = dspy.LM(self.config.lm, api_base=self.config.api_base, max_tokens=self.config.max_tokens)
        agent.set_lm(lm)
        self.agent = agent
        
    def forward(self, task: str) -> str:
        assert task, "Task cannot be empty"
        return self.agent(task=task)

# --- Main ---


def main():
    """Create AgentConfig with selected model."""
    model = os.getenv("MODEL")
    if model is None:
        model = select_model()

    # Add openrouter/ prefix if not already present
    if not model.startswith("openrouter/"):
        model = f"openrouter/{model}"
    
    config = AgentConfig()
    config.lm = model

    agent = AgentProgram(config)
    print(f"{BOLD}nanocode-dspy{RESET} | {DIM}{agent.config.lm} | {os.getcwd()}{RESET}\n")

    # Conversation history for context
    history = []

    while True:
        try:
            print(separator())
            user_input = input(f"{BOLD}{BLUE}❯{RESET} ").strip()
            print(separator())

            if not user_input:
                continue
            if user_input in ("/q", "exit"):
                break
            if user_input == "/c":
                history = []
                print(f"{GREEN}⏺ Cleared conversation{RESET}")
                continue

            # Build context from history
            context = f"Working directory: {os.getcwd()}\n"
            if history:
                context += "\nPrevious conversation:\n"
                for h in history[-5:]:  # Keep last 5 exchanges
                    context += f"User: {h['user']}\nAssistant: {h['assistant']}\n\n"

            task = f"{context}\nCurrent task: {user_input}"

            print(f"\n{CYAN}⏺{RESET} Thinking...", flush=True)

            # Run the ReAct agent
            result = agent(task=task)

            # Display the answer
            print(f"\n{CYAN}⏺{RESET} {render_markdown(result.answer)}")

            # Save to history
            history.append({"user": user_input, "assistant": result.answer})

            print()

        except (KeyboardInterrupt, EOFError):
            break
        except Exception as err:
            import traceback
            traceback.print_exc()
            print(f"{RED}⏺ Error: {err}{RESET}")


if __name__ == "__main__":
    main()