"""
Ollama CLI v3.0
License: MIT
Author: Gemini CLI
"""

import sys
import os
import json
import re
import logging
from typing import List, Dict, Optional, Tuple

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(os.path.expanduser("~/.ollama-cli.log")),
    ]
)
logger = logging.getLogger("ollama-cli")

# Add current directory to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from ollama_cli.core.config import load_config, save_config, update_config
from ollama_cli.core.ollama import OllamaClient
from ollama_cli.core.sessions import save_session, load_session, list_sessions
from ollama_cli.ui.formatter import console, print_markdown, print_error, print_status, StreamingDisplay
from ollama_cli.ui.repl import REPL
from ollama_cli.tools.base import registry
import ollama_cli.tools.filesystem
import ollama_cli.tools.system
import ollama_cli.tools.web
import ollama_cli.tools.git
import ollama_cli.tools.code
import ollama_cli.tools.knowledge
import ollama_cli.tools.memory
import ollama_cli.tools.execution
import ollama_cli.tools.media
from ollama_cli.integrations.mcp import connect_mcp_server, call_mcp_tool, MCP_TOOLS
from ollama_cli.integrations.lsp import start_lsp_server, lsp_goto_definition
from ollama_cli.integrations.notify import start_notify_thread, send_notification

VERSION = "3.0.0"

class OllamaCLI:
    def __init__(self):
        self.config = load_config()
        self.client = OllamaClient(self.config.get("ollama_url"))
        self.repl = REPL()
        self.messages = []
        self.available_models = []
        self.current_model = self.config.get("default_model", "llama3.2")
        self.system_prompt = self._build_system_prompt()
        self.max_tool_iterations = 5

    def _build_system_prompt(self) -> str:
        base_prompt = """You are a highly capable AI assistant powered by Ollama. 
You have access to a variety of tools to help you complete tasks.

CRITICAL RULES:
1. When you need to use a tool, you MUST use the EXACT format below.
2. Output ONLY the tool call. DO NOT explain what you are doing.
3. DO NOT repeat a tool call if you have already received a result for it.
4. If you have the answer, just give it. Do not use tools unless necessary.

<tool_call>
<tool_name>name_of_tool</tool_name>
<parameters>{"param": "value"}</parameters>
</tool_call>

You can call multiple tools in one response by repeating the block above."""
        
        tool_desc = registry.generate_system_prompt_snippet()
        
        # Add MCP tools if any
        if MCP_TOOLS:
            tool_desc += "\nMCP Tools:\n"
            for name, info in MCP_TOOLS.items():
                tool_desc += f"- {name}: {info['description']}\n"
        
        return f"{base_prompt}\n\n{tool_desc}"

    def reset_messages(self):
        self.messages = [{"role": "system", "content": self.system_prompt}]

    def handle_command(self, user_input: str) -> bool:
        """Handle slash commands."""
        if not user_input.startswith('/'):
            return False
            
        parts = user_input.split()
        cmd = parts[0].lower()
        
        try:
            if cmd == '/help':
                help_text = """
[bold cyan]Ollama CLI v3.0 Help[/bold cyan]

[bold]General Commands:[/bold]
  /help           - Show this help message
  /quit           - Exit the CLI
  /clear          - Clear conversation history

[bold]Model & Sessions:[/bold]
  /models         - List available models
  /model <name>   - Switch current model
  /sessions       - List saved sessions
  /save [name]    - Save current session
  /load <name>    - Load a saved session

[bold]Core Agent Tools:[/bold]
  [blue]Filesystem:[/blue] read_file, write_file, list_directory, grep_search, get_tree
  [blue]Code & Python:[/blue] run_python, replace_text, code_analyze_file, run_linter
  [blue]Web:[/blue] web_search, read_url
  [blue]Media:[/blue] generate_image, speak_text
  [blue]Knowledge:[/blue] kb_add, kb_search, remember_fact, recall_facts, clear_memory

[bold]Integrations:[/bold]
  /mcp connect <name> <cmd> [args] - Connect MCP server
  /lsp start <lang> [path]         - Start LSP server
  /notify setup <topic>            - Setup ntfy.sh remote control

[bold]Configuration:[/bold]
  /config ollama <url> - Set Ollama server URL (default: http://localhost:11434)
  /config comfy <url>  - Set ComfyUI server URL (default: http://127.0.0.1:8188)
  /config piper_path <path>  - Set path to Piper binary
  /config piper_model <path> - Set path to Piper voice model
"""
                console.print(help_text)
            elif cmd == '/config':
                if len(parts) == 1:
                    cfg = load_config()
                    config_display = f"""
[bold cyan]Current Configuration:[/bold cyan]
  [blue]Ollama URL:[/blue]  {cfg.get('ollama_url')}
  [blue]ComfyUI URL:[/blue] {cfg.get('comfy_url')}
  [blue]Image Model:[/blue] {cfg.get('image_model')}
  [blue]Piper Path:[/blue]  {cfg.get('piper', {}).get('path')}
  [blue]Piper Model:[/blue] {cfg.get('piper', {}).get('model')}
  [blue]Default Model:[/blue] {cfg.get('default_model')}
"""
                    console.print(config_display)
                elif len(parts) > 2:
                    subcmd = parts[1].lower()
                    val = parts[2]
                    if subcmd == 'ollama':
                        update_config("ollama_url", val)
                        print_status(f"Ollama URL updated to: {val}")
                        self.client = OllamaClient(val)
                    elif subcmd == 'comfy':
                        update_config("comfy_url", val)
                        print_status(f"ComfyUI URL updated to: {val}")
                    elif subcmd == 'piper_path':
                        cfg = load_config()
                        cfg["piper"]["path"] = val
                        save_config(cfg)
                        print_status(f"Piper binary path updated to: {val}")
                    elif subcmd == 'piper_model':
                        cfg = load_config()
                        cfg["piper"]["model"] = val
                        save_config(cfg)
                        print_status(f"Piper model path updated to: {val}")
                    else:
                        print_error("Usage: /config <ollama|comfy|piper_path|piper_model> <value>")
                else:
                    print_error("Usage: /config <ollama|comfy|piper_path|piper_model> <value>")
            elif cmd == '/quit':
                print_status("Goodbye!")
                sys.exit(0)
            elif cmd == '/clear':
                self.reset_messages()
                print_status("History cleared.")
            elif cmd == '/models':
                self.available_models = self.client.get_available_models()
                print_status(f"Available models: {', '.join(self.available_models)}")
            elif cmd == '/model':
                if len(parts) > 1:
                    self.current_model = parts[1]
                    update_config("default_model", self.current_model)
                    print_status(f"Switched to model: {self.current_model}")
                else:
                    print_status(f"Current model: {self.current_model}")
            elif cmd == '/sessions':
                sessions = list_sessions()
                print_status(f"Recent sessions: {', '.join(sessions[:10])}")
            elif cmd == '/save':
                name = parts[1] if len(parts) > 1 else None
                path = save_session(self.messages, name)
                print_status(f"Session saved to {path}")
            elif cmd == '/load':
                if len(parts) > 1:
                    messages = load_session(parts[1])
                    if messages:
                        self.messages = messages
                        print_status(f"Session '{parts[1]}' loaded.")
                    else:
                        print_error(f"Session '{parts[1]}' not found.")
                else:
                    print_error("Usage: /load <session_name>")
            elif cmd == '/mcp':
                if len(parts) > 3 and parts[1] == 'connect':
                    name = parts[2]
                    command = parts[3]
                    args = parts[4:] if len(parts) > 4 else None
                    result = connect_mcp_server(name, command, args)
                    print_status(result)
                    # Refresh system prompt with new tools
                    self.system_prompt = self._build_system_prompt()
                    self.reset_messages()
                else:
                    print_error("Usage: /mcp connect <name> <command> [args...]")
            elif cmd == '/lsp':
                if len(parts) > 2 and parts[1] == 'start':
                    lang = parts[2]
                    path = parts[3] if len(parts) > 3 else "."
                    result = start_lsp_server(lang, path)
                    print_status(result)
                else:
                    print_error("Usage: /lsp start <language> [path]")
            elif cmd == '/notify':
                if len(parts) > 2 and parts[1] == 'setup':
                    topic = parts[2]
                    self.config["ntfy"]["enabled"] = True
                    self.config["ntfy"]["topic"] = topic
                    save_config(self.config)
                    start_notify_thread(self.config, self.process_remote_message)
                    print_status(f"Notifications enabled for topic: {topic}")
                else:
                    print_error("Usage: /notify setup <topic>")
            else:
                print_error(f"Unknown command: {cmd}")
        except Exception as e:
            print_error(f"Command error: {e}")
            
        return True

    def process_remote_message(self, message: str) -> str:
        """Handle messages coming from ntfy.sh"""
        temp_messages = self.messages + [{"role": "user", "content": message}]
        try:
            response = ""
            for chunk in self.client.chat(temp_messages, self.current_model, stream=False):
                response += chunk
            return response
        except Exception as e:
            return f"Error: {e}"

    def run(self):
        console.print(f"[bold cyan]Ollama CLI v{VERSION}[/bold cyan]")
        
        # Initial model detection
        try:
            self.available_models = self.client.get_available_models()
            if not self.available_models:
                print_error("No models found. Make sure Ollama is running.")
            else:
                if self.current_model not in self.available_models:
                    self.current_model = self.available_models[0]
                print_status(f"Using model: [bold green]{self.current_model}[/bold green]")
        except Exception as e:
            print_error(f"Could not connect to Ollama: {e}")

        self.reset_messages()
        
        # Start notification listener if enabled
        if self.config.get("ntfy", {}).get("enabled"):
            start_notify_thread(self.config, self.process_remote_message)
            print_status(f"Notification listener active on topic: {self.config['ntfy']['topic']}")

        while True:
            user_input = self.repl.get_input()
            
            if not user_input:
                continue
                
            if self.handle_command(user_input):
                continue
                
            self.messages.append({"role": "user", "content": user_input})
            self.process_chat_cycle()

    def process_chat_cycle(self):
        """The main loop for chat and tool execution."""
        executed_tool_calls = set()
        for iteration in range(self.max_tool_iterations):
            full_response = ""
            try:
                print_status(f"Thinking ([dim]{self.current_model}[/dim])...")
                with StreamingDisplay() as display:
                    for chunk in self.client.chat(self.messages, self.current_model):
                        display.update(chunk)
                        full_response += chunk
                
                self.messages.append({"role": "assistant", "content": full_response})
                
                # Check for tool calls
                tool_calls = self.parse_tool_calls(full_response)
                if not tool_calls:
                    break
                
                # Execute tool calls
                results = []
                for name, params in tool_calls:
                    # Prevent identical tool call loops across the entire cycle
                    call_key = f"{name}:{json.dumps(params, sort_keys=True)}"
                    if call_key in executed_tool_calls:
                        print_status(f"Skipping duplicate tool call: [dim]{name}[/dim]")
                        continue
                    executed_tool_calls.add(call_key)
                    
                    result = self.execute_tool(name, params)
                    # Show result to user
                    print_status(f"Result: [dim]{result}[/dim]")
                    results.append(f"Tool Result ({name}):\n{result}")
                
                if not results:
                    # If we had tool calls but they were all duplicates, we must stop
                    break
                    
                # Add all results to conversation
                self.messages.append({"role": "user", "content": "\n\n".join(results) + "\n\nTool execution complete. If the task is finished, provide a final response to the user. Do not repeat the tool call."})
                
                # If we made it here, loop again to let the model process results
                
            except Exception as e:
                print_error(str(e))
                logger.exception("Chat cycle error")
                break

    def parse_tool_calls(self, text: str) -> List[tuple]:
        calls = []
        
        # 1. Try to find <tool_call> blocks (handles unclosed tags)
        tc_blocks = re.findall(r'<tool_call>(.*?)(?:</tool_call>|$)', text, re.DOTALL)
        if tc_blocks:
            for block in tc_blocks:
                if not block.strip(): continue
                # Try to find tool_name tag (handles unclosed)
                name_match = re.search(r'<tool_name>(.*?)(?:</tool_name>|$)', block, re.DOTALL)
                name = name_match.group(1).strip() if name_match else None
                
                # Try to find parameters tag (handles unclosed)
                params_match = re.search(r'<parameters>(.*?)(?:</parameters>|$)', block, re.DOTALL)
                if name and params_match:
                    try:
                        params_str = params_match.group(1).strip()
                        # Auto-fix unclosed JSON if needed
                        if params_str.startswith('{') and not params_str.endswith('}'):
                            params_str += '}'
                        calls.append((name, json.loads(params_str)))
                    except:
                        # Last ditch: search for any JSON-like object in the parameters string
                        json_match = re.search(r'(\{.*\})', params_str, re.DOTALL)
                        if json_match:
                            try: calls.append((name, json.loads(json_match.group(1))))
                            except: pass
        
        # 2. If no <tool_call> blocks, try searching for tool-specific tags
        if not calls:
            for tool_name in registry.tools.keys():
                # Handle <tool_name ... /> or <tool_name ... >
                attr_pattern = fr'<{tool_name}\s+(.*?)[\s/]*>'
                attr_matches = re.findall(attr_pattern, text, re.DOTALL)
                for attr_str in attr_matches:
                    attrs = {}
                    kv_pairs = re.findall(r'(\w+)\s*=\s*["\'](.*?)["\']', attr_str)
                    for k, v in kv_pairs:
                        attrs[k] = v
                    if attrs:
                        calls.append((tool_name, attrs))
                
                # Handle <tool_name> {json} </tool_name>
                block_pattern = fr'<{tool_name}>(.*?)(?:</{tool_name}>|$)'
                block_matches = re.findall(block_pattern, text, re.DOTALL)
                for b in block_matches:
                    json_match = re.search(r'(\{.*\})', b, re.DOTALL)
                    if json_match:
                        try: calls.append((tool_name, json.loads(json_match.group(1))))
                        except: pass
        
        # 3. Final fallback: look for tool_name followed by JSON { ... }
        if not calls:
            for tool_name in registry.tools.keys():
                json_pattern = fr'{tool_name}\s*({{.*?}})'
                json_matches = re.findall(json_pattern, text, re.DOTALL)
                for json_str in json_matches:
                    try:
                        calls.append((tool_name, json.loads(json_str.strip())))
                    except:
                        pass

        return calls

    def execute_tool(self, name: str, params: Dict) -> str:
        """Execute a tool (standard or MCP)."""
        # Try MCP tools first
        if name in MCP_TOOLS:
            mcp_tool = MCP_TOOLS[name]
            print_status(f"Executing MCP tool [bold blue]{mcp_tool['name']}[/bold blue]...")
            return call_mcp_tool(mcp_tool["server"], mcp_tool["name"], params)
        
        # Try standard tools
        tool = registry.get_tool(name)
        if tool:
            print_status(f"Executing [bold blue]{name}[/bold blue]...")
            try:
                result = tool.execute(**params)
                # Notification for successful file/system ops
                if name == "write_file" and "Successfully" in str(result):
                    send_notification(self.config, "File Created", f"Path: {params.get('path')}")
                return result
            except Exception as e:
                return f"Error executing tool: {e}"
        
        return f"Unknown tool: {name}"

def main():
    try:
        cli = OllamaCLI()
        cli.run()
    except KeyboardInterrupt:
        print_status("Interrupted by user. Exiting.")
        sys.exit(0)
    except Exception as e:
        print_error(f"Fatal error: {e}")
        logger.exception("Fatal error")
        sys.exit(1)

if __name__ == "__main__":
    main()