"""
Ollama Agent with Tool Calling Support

Questo modulo implementa un agente che supporta tool calling per Ollama
usando un approccio di parsing delle risposte invece di bind_tools.
"""
import json
import re
import uuid
import subprocess
from typing import List, Dict, Any, Optional, Callable
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path

from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage
from langchain_core.tools import tool, Tool


# =============================================================================
# TOOL CALL PATTERNS
# =============================================================================

# Old format: [TOOL_CALL]...[/TOOL_CALL]
TOOL_CALL_PATTERN_OLD = r'\[TOOL_CALL\]\s*(.*?)\s*\[\/TOOL_CALL\]'

# New format: TOOL:toolname|arg1=val1|arg2=val2
TOOL_CALL_PATTERN_NEW = r'^TOOL:(\w+)(?:\|(.*))?$'

TOOL_RESULT_PATTERN = r'\[TOOL_RESULT\]\s*(.*?)\s*\[\/TOOL_RESULT\]'


# =============================================================================
# BASIC BUILT-IN TOOLS
# =============================================================================

def get_builtin_tools() -> List[Tool]:
    """
    Restituisce i tool built-in disponibili per l'agente.
    """
    tools = []
    
    # Tool: Info sistema - usa @tool decorator
    @tool
    def system_info() -> str:
        """Get system information (OS, Python version, current directory)."""
        import platform
        import os
        info = []
        info.append(f"OS: {platform.system()} {platform.release()}")
        info.append(f"Python: {platform.python_version()}")
        info.append(f"Current dir: {os.getcwd()}")
        info.append(f"User: {os.getenv('USER', 'unknown')}")
        return "\n".join(info)
    
    tools.append(system_info)
    
    # Tool: Lista file
    @tool
    def ls(path: str) -> str:
        """List files in a directory."""
        try:
            dir_path = Path(path).expanduser()
            if not dir_path.exists():
                return f"Error: Directory '{path}' not found"
            if not dir_path.is_dir():
                return f"Error: '{path}' is not a directory"
            
            items = []
            for item in sorted(dir_path.iterdir()):
                item_type = "DIR" if item.is_dir() else "FILE"
                items.append(f"{item_type}: {item.name}")
            return "\n".join(items) if items else "Directory is empty"
        except Exception as e:
            return f"Error listing directory: {str(e)}"
    
    tools.append(ls)
    
    # Tool: Esegui comando shell
    @tool
    def execute(command: str) -> str:
        """Execute a shell command. Use this for file operations, running scripts, git commands, etc."""
        try:
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=30
            )
            output = []
            if result.stdout:
                output.append(f"STDOUT:\n{result.stdout}")
            if result.stderr:
                output.append(f"STDERR:\n{result.stderr}")
            if result.returncode != 0:
                output.append(f"Exit code: {result.returncode}")
            return "\n".join(output) if output else "Command executed with no output"
        except subprocess.TimeoutExpired:
            return "Error: Command timed out after 30 seconds"
        except Exception as e:
            return f"Error: {str(e)}"
    
    tools.append(execute)
    
    # Tool: Leggi file
    @tool
    def read_file(path: str) -> str:
        """Read the contents of a file."""
        try:
            file_path = Path(path).expanduser()
            if not file_path.exists():
                return f"Error: File '{path}' not found"
            if file_path.is_dir():
                return f"Error: '{path}' is a directory, not a file"
            return file_path.read_text(encoding='utf-8')
        except Exception as e:
            return f"Error reading file: {str(e)}"
    
    tools.append(read_file)
    
    # Tool: Scrivi file
    @tool
    def write_file(path: str, content: str) -> str:
        """Write content to a file. WARNING: overwrites existing files!"""
        try:
            file_path = Path(path).expanduser()
            file_path.parent.mkdir(parents=True, exist_ok=True)
            file_path.write_text(content, encoding='utf-8')
            return f"Successfully wrote to {path}"
        except Exception as e:
            return f"Error writing file: {str(e)}"
    
    tools.append(write_file)
    
    # Tool: Cerca file
    @tool
    def search(directory: str, pattern: str) -> str:
        """Search for files by name pattern."""
        try:
            dir_path = Path(directory).expanduser()
            if not dir_path.exists():
                return f"Error: Directory '{directory}' not found"
            
            matches = list(dir_path.rglob(pattern))
            if not matches:
                return f"No files matching '{pattern}' found in {directory}"
            
            results = []
            for match in matches[:50]:
                results.append(str(match))
            
            return f"Found {len(matches)} matches (showing first 50):\n" + "\n".join(results)
        except Exception as e:
            return f"Error searching files: {str(e)}"
    
    tools.append(search)
    
    return tools


@dataclass
class ToolCall:
    """Rappresenta una chiamata a tool"""
    id: str
    name: str
    arguments: Dict[str, Any]


class OllamaToolAgent:
    """
    Agente che supporta tool calling per Ollama e provider simili.
    
    Strategia: Il modello risponde con marcatori speciali per le chiamate tool.
    Noi parsiamo, eseguiamo, e riportiamo i risultati.
    """
    
    def __init__(
        self,
        llm,
        tools: List[Tool],
        system_prompt: Optional[str] = None,
        max_iterations: int = 10
    ):
        self.llm = llm
        self.tools = {tool.name: tool for tool in tools}
        self.max_iterations = max_iterations
        
        # Costruisci il system prompt con la lista dei tool
        if system_prompt is None:
            system_prompt = self._default_system_prompt()
        else:
            system_prompt = system_prompt + "\n\n" + self._tool_instructions()
        
        self.messages: List[BaseMessage] = [SystemMessage(content=system_prompt)]
        
        # Tool per la memoria
        self.short_term_memory: List[Dict] = []
        
    def _default_system_prompt(self) -> str:
        """System prompt di default"""
        return """You are a helpful AI assistant. When users ask about:
- system info, OS, operating system, what computer: respond with TOOL:system_info
- list files, directory contents, ls: respond with TOOL:ls|path=.

Example: User asks "what OS am I using?" -> You respond: TOOL:system_info
Example: User asks "list files" -> You respond: TOOL:ls|path=.

For other questions, just answer normally without tools."""
    
    def _tool_instructions(self) -> str:
        """Genera le istruzioni per i tool disponibili"""
        instructions = []
        instructions.append("\n\n# AVAILABLE TOOLS\n")
        instructions.append("You have access to the following tools:\n\n")
        
        for name, tool in self.tools.items():
            # Estrai la descrizione dal docstring del tool
            desc = tool.description or "No description"
            # Pulisci la descrizione da caratteri speciali
            desc = desc.replace('"', '\\"').replace('\n', ' ')
            instructions.append(f"## {name}\n")
            instructions.append(f"Description: {desc}\n")
            
            # Parsa i parametri se presenti
            if hasattr(tool, 'args_schema') and tool.args_schema:
                instructions.append("Parameters:\n")
                schema = tool.args_schema
                # Handle both dict and BaseModel schemas
                try:
                    model_json_schema = getattr(schema, 'model_json_schema', None)  # type: ignore
                    if callable(model_json_schema):
                        schema = model_json_schema()
                except Exception:
                    pass
                if isinstance(schema, dict) and 'properties' in schema:
                    for param_name, param_info in schema['properties'].items():
                        param_type = param_info.get('type', 'any')
                        param_desc = param_info.get('description', '')
                        required = param_name in schema.get('required', [])
                        req_marker = "(required)" if required else "(optional)"
                        instructions.append(f"  - {param_name}: {param_type} {req_marker} - {param_desc}\n")
            
            instructions.append("\n")
        
        instructions.append("\n# IMPORTANT RULES:\n")
        instructions.append("1. When you need to use a tool, respond ONLY with [TOOL_CALL] blocks\n")
        instructions.append("2. After receiving the result, continue your response\n")
        instructions.append("3. For simple questions, respond directly without tools\n")
        instructions.append("4. Use valid JSON in tool calls\n")
        instructions.append("5. If a tool fails, explain the error to the user\n")
        
        return ''.join(instructions)
    
    def _parse_tool_calls(self, text: str) -> List[ToolCall]:
        """Parse le chiamate tool dalla risposta del modello"""
        tool_calls = []
        
        # Try new format first: TOOL:toolname|arg1=val1
        lines = text.strip().split('\n')
        for line in lines:
            line = line.strip()
            if line.startswith('TOOL:'):
                # Parse: TOOL:toolname|arg1=val1|arg2=val2
                match = re.match(r'^TOOL:(\w+)(?:\|(.+))?$', line)
                if match:
                    tool_name = match.group(1)
                    if tool_name in self.tools:
                        args = {}
                        if match.group(2):
                            # Parse arguments: arg1=val1|arg2=val2
                            for arg_part in match.group(2).split('|'):
                                if '=' in arg_part:
                                    key, val = arg_part.split('=', 1)
                                    args[key] = val
                        tool_calls.append(ToolCall(
                            id=str(uuid.uuid4())[:8],
                            name=tool_name,
                            arguments=args
                        ))
                        # If we found a valid tool call, return only that one
                        return tool_calls
        
        # Try old format: [TOOL_CALL]{...}[/TOOL_CALL]
        matches = re.finditer(TOOL_CALL_PATTERN_OLD, text, re.DOTALL)
        for match in matches:
            try:
                data = json.loads(match.group(1))
                tool_name = data.get('tool')
                if tool_name and tool_name in self.tools:
                    tool_calls.append(ToolCall(
                        id=str(uuid.uuid4())[:8],
                        name=tool_name,
                        arguments=data.get('arguments', {})
                    ))
            except (json.JSONDecodeError, TypeError):
                continue
        
        return tool_calls
    
    def _execute_tool(self, tool_call: ToolCall) -> Any:
        """Esegui un singolo tool"""
        if tool_call.name not in self.tools:
            available = ", ".join(self.tools.keys())
            return {
                "error": f"Tool '{tool_call.name}' not found",
                "available_tools": available
            }
        
        try:
            tool = self.tools[tool_call.name]
            
            # Per tool senza argomenti, non passare arguments
            if tool_call.arguments:
                result = tool.invoke(tool_call.arguments)
            else:
                result = tool.invoke({})
            
            # Converti result in stringa per il modello
            if result is None:
                return {"result": "Tool executed successfully"}
            elif isinstance(result, (dict, list)):
                return {"result": json.dumps(result, ensure_ascii=False, indent=2)}
            else:
                return {"result": str(result)}
                
        except Exception as e:
            return {"error": str(e)}
    
    def _format_tool_result(self, tool_call: ToolCall, result: Any) -> str:
        """Formatta il risultato di un tool per il modello"""
        return f"\n[TOOL_RESULT]\nTool: {tool_call.name}\nResult: {json.dumps(result, ensure_ascii=False)}\n[/TOOL_RESULT]\n"
    
    def _detect_intent(self, text: str) -> Optional[ToolCall]:
        """Detect intents using both patterns and LLM interpretation"""
        import re
        
        text_lower = text.lower().strip()
        
        # Extract potential file paths (case insensitive)
        path_match = re.search(r'([/\w.-]+\.\w+)', text, re.IGNORECASE)
        file_path = path_match.group(1) if path_match else None
        
        # Execute command patterns - check FIRST for shell commands with flags
        # Shell commands like "ls -lha", "git status", "pwd", "echo" should use execute
        shell_commands = ['ls', 'git', 'cat', 'grep', 'find', 'pwd', 'echo', 'cd', 
                         'mkdir', 'rm', 'cp', 'mv', 'chmod', 'head', 'tail', 'sort', 'uniq']
        # Check for command alone or with arguments
        for cmd in shell_commands:
            if text_lower.strip() == cmd or text_lower.startswith(cmd + ' '):
                return ToolCall(id='intent', name='execute', arguments={'command': text.strip()})
        
        # System info patterns (more comprehensive)
        system_patterns = [
            'sistema operativo', 'so sto', 'che os', 'che sistema', 'che so',
            'sto usando', 'che computer', 'che tipo di sistema',
            'versione python', 'che versione', 'dettagli sistema',
            'operating system', 'what os', 'which system', 'what system',
            'which computer', 'system info', 'python version'
        ]
        if any(p in text_lower for p in system_patterns):
            return ToolCall(id='intent', name='system_info', arguments={})
        
        # List files patterns
        list_patterns = [
            'lista', 'listare', 'listami', 'elenco', 'mostrami i file', 'mostrami',
            'file nella cartella', 'file del', 'cartella corrente', 'mostra i file',
            'cosa c\'è', 'cosa c\'e', 'cosa contiene', 'cosa trovo',
            'list files', 'list directory', 'ls', 'show files', 'show me',
            'what\'s in', 'what is in', 'files here', 'directory contents'
        ]
        if any(p in text_lower for p in list_patterns):
            return ToolCall(id='intent', name='ls', arguments={'path': '.'})
        
        # Search patterns
        search_patterns = [
            'cerca', 'cercare', 'trova', 'trovare', 'dove si trova',
            'search', 'find', 'file', '*.py', '*.md', '*.txt'
        ]
        # If the message mentions file patterns, use search tool
        if 'cerca' in text_lower or 'search' in text_lower or 'find' in text_lower:
            # Extract directory and pattern
            dir_match = re.search(r'(?:in|alla|nella)\s+([/\w.-]+)', text)
            pattern_match = re.search(r'\*\.(\w+)', text)
            directory = dir_match.group(1) if dir_match else '.'
            pattern = f"*.{pattern_match.group(1)}" if pattern_match else "*"
            return ToolCall(id='intent', name='search', arguments={'directory': directory, 'pattern': pattern})
        
        # Read/show file patterns - check for ANY read keyword AND a file path
        read_keywords = ['leggi', 'leggere', 'mostrami', 'apri', 'aprimi', 'visualizza', 
                        'vediamo', 'vedi', 'show me', 'show the', 'read', 'open', 
                        'display', 'view', 'contenuto', 'che dice']
        has_read_keyword = any(k in text_lower for k in read_keywords)
        
        if has_read_keyword and file_path:
            return ToolCall(id='intent', name='read_file', arguments={'path': file_path})
        
        # Write file patterns
        write_keywords = ['scrivi', 'scrivere', 'crea', 'creare', 'salva', 'salvare',
                        'write', 'create', 'save', 'make file']
        has_write_keyword = any(k in text_lower for k in write_keywords)
        
        if has_write_keyword and file_path:
            # Extract content - try multiple patterns
            content = None
            
            # Pattern 1: "con dentro 'testo'" o "con 'testo'"
            content_match = re.search(r'con\s+(?:dentro\s+)?["\']([^"\']+)["\']', text, re.IGNORECASE)
            if content_match:
                content = content_match.group(1)
            
            # Pattern 2: testo tra virgolette anywhere
            if not content:
                content_match = re.search(r'["\']([^"\']+)["\']', text)
                if content_match:
                    content = content_match.group(1)
            
            # Pattern 3: contenuto dopo "con" fino alla fine (senza virgolette)
            if not content:
                # Match: con <resto della frase senza virgolette iniziali/finali>
                after_con = re.search(r'con\s+["\']?(.+?)["\']?\s*$', text, re.IGNORECASE | re.DOTALL)
                if after_con:
                    content = after_con.group(1).strip()
            
            if not content:
                content = "Content not specified"
            
            return ToolCall(id='intent', name='write_file', arguments={'path': file_path, 'content': content})
        
        exec_patterns = [
            'esegui', 'eseguire', 'run', 'run command', 'lanciare',
            'avvia', 'comando', 'terminale', 'bash', 'shell'
        ]
        if any(p in text_lower for p in exec_patterns):
            cmd_match = re.search(r'[`"]?([a-zA-Z0-9_.\s-]+)[`"]?\s*$', text)
            if cmd_match:
                return ToolCall(id='intent', name='execute', arguments={'command': cmd_match.group(1).strip()})
        
        return None
    
    def invoke(self, input_dict: Dict) -> Dict:
        """
        Invoca l'agente con un input.
        
        L'agente può fare multiple iterazioni tool-call prima di rispondere.
        """
        # Estrai il messaggio utente
        user_message = input_dict.get("messages", [])[-1]
        
        # Gestisci diversi formati di messaggio
        if isinstance(user_message, dict):
            user_content = user_message.get("content", str(user_message))
        elif hasattr(user_message, 'content'):
            user_content = user_message.content
        else:
            user_content = str(user_message)
        
        # Aggiungi il messaggio utente alla cronologia
        self.messages.append(HumanMessage(content=user_content))
        
        # Try to detect intent directly first
        detected_tool = self._detect_intent(user_content)
        if detected_tool:
            result = self._execute_tool(detected_tool)
            
            # Format result as a clean response
            response_text = ""
            try:
                if isinstance(result, str):
                    try:
                        result = json.loads(result)
                    except:
                        response_text = result
                
                if isinstance(result, dict):
                    if "result" in result:
                        response_text = result["result"]
                    elif "error" in result:
                        response_text = f"Errore: {result['error']}"
                    else:
                        response_text = json.dumps(result, indent=2)
            except Exception as e:
                response_text = str(result) if not response_text else response_text
            
            return {"messages": self.messages, "response": response_text}
        
        # Se nessun intent detectato, usa il modello per rispondere normalmente
        # (senza tool calling per ora - i modelli locali non lo supportano bene)
        response = self.llm.invoke(self.messages)
        self.messages.append(response)
        
        response_text = response.content if hasattr(response, 'content') else str(response)
        return {"messages": self.messages, "response": response_text}
    
    def clear_history(self):
        """Pulisci la cronologia ma mantieni il system prompt"""
        system_msg = self.messages[0] if self.messages and isinstance(self.messages[0], SystemMessage) else None
        self.messages = [system_msg] if system_msg else []


def create_ollama_tool_agent(
    model,
    tools: Optional[List[Tool]] = None,
    system_prompt: Optional[str] = None,
    include_builtin_tools: bool = True
) -> OllamaToolAgent:
    """
    Factory function per creare un OllamaToolAgent.
    
    Args:
        model: Il modello LLM da usare
        tools: Lista di tool personalizzati (opzionale)
        system_prompt: System prompt personalizzato (opzionale)
        include_builtin_tools: Se True, include i tool built-in (execute, read_file, etc.)
    """
    if tools is None:
        tools = []
    
    # Aggiungi i tool built-in
    if include_builtin_tools:
        builtin = get_builtin_tools()
        tools = builtin + tools
    
    return OllamaToolAgent(
        llm=model,
        tools=tools,
        system_prompt=system_prompt
    )