#!/usr/bin/env python3

#!/usr/bin/env python3
# file manipulation in 'workspace'

import sys
import os
import json
import time
from pathlib import Path
import requests

# ===========================
# configuration
# ===========================
API_KEY = 'XXX'
API_URL = 'https://litellm.s.studiumdigitale.uni-frankfurt.de/v1/chat/completions'
MODEL   = 'openai-gpt-oss-120b'

BASE_DIR = Path(__file__).parent.resolve()
HISTORY_FILE = BASE_DIR / 'chat_history.jsonl'
WORKSPACE = BASE_DIR / 'workspace'

# Ensure workspace exists
WORKSPACE.mkdir(parents=True, exist_ok=True)

# ===========================
# Path handling
# ===========================
def safe_path(path_str: str = '') -> Path:
    """
    Resolves a relative path within the workspace sandbox safely,
    preventing directory traversal attacks.
    """
    root = WORKSPACE.resolve()
    path_str = path_str.strip().lstrip('/')
    
    if not path_str:
        return root
        
    candidate = (root / path_str).resolve()
    
    # Check if the path resides inside the workspace root
    if root not in candidate.parents and candidate != root:
        raise PermissionError('Access denied')
        
    return candidate

# ===========================
# Tool storage
# ===========================
TOOL_REGISTRY = {}       # all tool definitions
TOOL_HISTORY  = []       # of current session

# ===========================
# Tool: listing files
# ===========================
TOOL_REGISTRY["list_files"] = {\
  'type'    : 'function',
  'function': { 'name': 'list_files',
                'description': 'Recursively list files and directories in the workspace. Returns a formatted list showing [DIR] for directories and [FILE] for files.',
          'parameters': { 'type': 'object',
                'properties': { 'path': { 'type': 'string',
                'description': 'Optional subdirectory path. Leave empty for workspace root.'
                                              }
                                  }
                        }
              }
                               }

def list_files(path: str = '') -> str:
    try:
        target_dir = safe_path(path)
        if not target_dir.is_dir():
            return 'ERROR: Not a directory'
            
        root = WORKSPACE.resolve()
        result = []
        
# Recursive globbing to match PHP's RecursiveIteratorIterator
        for item in target_dir.rglob('*'):
            # Calculate relative path from workspace root
            relative = item.relative_to(root)
            if item.is_dir():
                result.append(f'[DIR]  {relative}')
            else:
                result.append(f'[FILE] {relative}')
                
# Natural/Case-insensitive sorting approximation
        result.sort(key=lambda s: s.lower())
        
        if not result:
            return '(empty workspace)'
            
        return '\n'.join(result)
    except Exception as e:
        return f'ERROR: {str(e)}'

# ===========================
# Tool: read file
# ===========================
TOOL_REGISTRY["read_file"]  = {\
  'type'    : 'function',
  'function': { 'name': 'read_file',
                'description': 'Read the contents of a text file. Returns the full content or a truncated version for large files.',
          'parameters': { 'type': 'object',
                    'properties': { 'path': { 'type': 'string',
                               'description': 'Path to the file to read.'
                                              }
                                  },
                      'required': ['path']
                        }
              }
                               }

def read_file(path: str) -> str:
    try:
        file_path = safe_path(path)
        if not file_path.exists():
            return 'ERROR: File not found'
        if not file_path.is_file():
            return 'ERROR: Not a file'
            
        content = file_path.read_text(encoding='utf-8', errors='replace')
        
        if len(content) > 100000:
            return content[:100000] + "\n\n[TRUNCATED]"
            
        return content
    except Exception as e:
        return f'ERROR: {str(e)}'

# ===========================
# Tool: write file
# ===========================
TOOL_REGISTRY["write_file"] = {
            'type': 'function',
            'function': {
                'name': 'write_file',
                'description': 'Write content to a file. Creates the file if it does not exist, or overwrites it if it does.',
                'parameters': {
                    'type': 'object',
                    'properties': {
                        'path': {'type': 'string', 'description': 'Path where the file should be written'},
                        'content': {'type': 'string', 'description': 'Content to write to the file'}
                    },
                    'required': ['path', 'content']
                }
            }
        }

def write_file(path: str, content: str) -> str:
    try:
        file_path = safe_path(path)

        # Ensure parent directories exist
        file_path.parent.mkdir(parents=True, exist_ok=True)
        file_path.write_text(content, encoding='utf-8')
        return f'{"Created" if not file_path.exists() else "Overwrote"} file: {path}'
    except Exception as e:
        return f'ERROR: {str(e)}'

# ===========================
# Tool: append to file
# ===========================
TOOL_REGISTRY["append_file"]  = {\
            'type': 'function',
            'function': {
                'name': 'append_file',
                'description': 'Append text to an existing file. Creates the file if it does not exist.',
                'parameters': { 
                    'type': 'object',
                    'properties': {
                        'path': {'type': 'string', 'description': 'Path to the file to append to'},
                        'content': {'type': 'string', 'description': 'Content to append to the file'}
                    },
                    'required': ['path', 'content']
                }
            }
        }

def append_file(path: str, content: str) -> str:
    try:
        file_path = safe_path(path)
        # Ensure parent directories exist
        file_path.parent.mkdir(parents=True, exist_ok=True)
        
        with open(file_path, 'a', encoding='utf-8') as f:
            f.write(content)
            
        return f'Appended to file: {path}'
    except Exception as e:
        return f'ERROR: {str(e)}'

# ===========================
# Tool-handlers (dispatch table pattern)
# Mapping tool names (strings) to handler functions (lambdas)
# Avoids exposing other functions to the LLM
# ===========================
TOOL_HANDLERS = {
    'list_files':   lambda args: list_files(args.get('path', '')),
    'read_file':    lambda args: read_file(args.get('path')),
    'write_file':   lambda args: write_file(args.get('path'), 
                           args.get('content')),
    'append_file':  lambda args: append_file(args.get('path'), 
                           args.get('content')),
                }

# ===========================
# chat history: loading
# ===========================
def load_history(history_path: Path) -> list:
    if not history_path.exists():
        return []
        
    messages = []
    try:
        with open(history_path, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    messages.append(json.loads(line))
                except json.JSONDecodeError:
                    continue
    except Exception:
        return []
    return messages

# ===========================
# chat history: saving 
# ===========================
def save_history_message(history_path: Path, role: str, 
                         content: str) -> None:
    entry = {'role': role, 'content': content}
    with open(history_path, 'a', encoding='utf-8') as f:
        f.write(json.dumps(entry, ensure_ascii=False) + '\n')

# ===========================
# model call: retrying
# ===========================
def call_model_with_retry(api_url, api_key, model, messages, max_retries=3):
    for i in range(max_retries):
        try:
            return call_model(api_url, api_key, model, messages)
        except Exception as e:
            if i == max_retries - 1:
                raise e
            time.sleep(1)

# ===========================
# model call: request  
# ===========================
def call_model(api_url: str, api_key: str, model: str, 
               messages: list) -> dict:
    tools = [TOOL_REGISTRY["list_files"],
             TOOL_REGISTRY["read_file"],
             TOOL_REGISTRY["write_file"],
             TOOL_REGISTRY["append_file"] ]

    payload = {
        'model': model,
        'messages': messages,
        'tools': tools,
        'tool_choice': 'auto',
        'temperature': 0.2
    }

    headers = {'Content-Type': 'application/json'}
    if api_key:
        headers['Authorization'] = f'Bearer {api_key}'

    try:
        response = requests.post(api_url, json=payload, headers=headers, timeout=60)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        # Debugging fallback matching PHP implementation
        print("\nDEBUG: Request failed or invalid response.")
        if 'response' in locals() and response is not None:
            print(f"DEBUG: Raw response (first 500 chars):\n{response.text[:500]}")
        raise Exception(f"API Error: {str(e)}")

# ===========================
# Agent loop
# ===========================
def run_agent(api_url: str, api_key: str, model: str, 
              messages: list) -> str:
    """looping as a long a the LLM call new tools"""
    while True:
        response = call_model_with_retry(api_url, api_key, model, messages)

        if 'choices' not in response or not response['choices']:
            print("\nDEBUG: Response structure:")
            print(response.keys())
            raise Exception('Unexpected API response structure')

        message = response['choices'][0]['message']
        tool_calls = message.get('tool_calls')

        if not tool_calls:        # LLM finished; exit run_agent()
            return message.get('content') or ''

# Append assistant response containing tool calls to conversation history
        messages.append(message)

        for tool_call in tool_calls:
            name = tool_call['function']['name']
#           print("# in run_agent() tool name : ", name, "  \n")
            TOOL_HISTORY.append(name)                # update history
            try:
                args = json.loads(tool_call['function']['arguments'])
            except (json.JSONDecodeError, TypeError):
                args = {}

            try:
              handler = TOOL_HANDLERS.get(name)      # call this tool
              result = handler(args) if handler else 'ERROR: Unknown tool'
            except Exception as e:
              result = f'ERROR: {str(e)}'

# Format the tool execution result role back into the array context
            messages.append({'role': 'tool',
                     'tool_call_id': tool_call['id'],
                             'name': name,
                          'content': result})

# ===========================
# System prompt
# ===========================
def get_sys_prompt():
  available_tools = [tool['function']['name']\
                 for tool in TOOL_REGISTRY.values()]
  SP = {
    'role': 'system',
    'content': (
        '# Workspace File Assistant\n\n'
        'You are an automated file operator. Your job is to execute file operations using the available tools.\n\n'
        '## Critical Rules\n'
        '- **MUST** use tools for all file operations\n'
        '- **NEVER** invent file contents\n'
        '- **ALWAYS** read before modifying\n'
        '- **NEVER** answer from prior knowledge\n'
        '- **NEVER** perform operations outside the given workspace\n\n'
        '## Available Tools\n'
        f'- {", ".join(available_tools)}\n\n'
        '## Paths\n'
        '- All paths are relative to the workspace root.'
               )
       }
  print("# ===========")
  print(SP['content'])
  print("# ===========\n")
  return SP

# ===========================
# Main
# ===========================
def main():
    print(f"\nMinimal File Agent")
    print(f"Workspace: {WORKSPACE}")
    print("Type 'exit' to quit\n")

    with open(HISTORY_FILE, "w") as f:   # empty hitory file
       pass
    history = load_history(HISTORY_FILE) # start fresh

    system_prompt = get_sys_prompt()
    old_nTools = 0

    while True:
        try:
            nTools  = len(TOOL_HISTORY)
            cT      = TOOL_HISTORY[old_nTools:nTools] if nTools>0 else ""
            old_nTools = nTools
#           print("# old_nTools ", old_nTools)
#           print("#     nTools ",     nTools)
            print(f"\r\033[K[{nTools}] {cT} ", end="", flush=True)
            user_input = input('> ').strip()
#           print(user_input)
        except (KeyboardInterrupt, EOFError):
            print("\nBye!")
            sys.exit(0)

        if not user_input:
            continue

        if user_input.lower() in ['exit', 'quit']:
            print("Bye!")
            sys.exit(0)

        save_history_message(HISTORY_FILE, 'user', user_input)

# Build context list matching structural constraints of the PHP script
        messages = [system_prompt] + history + [{'role': 'user', 'content': user_input}]

        try:
            reply = run_agent(API_URL, API_KEY, MODEL, messages)
            
            print(f"\n{reply}\n")
            save_history_message(HISTORY_FILE, 'assistant', reply)

# Keep runtime user history synchronized
            history.append({'role': 'user', 'content': user_input})
            history.append({'role': 'assistant', 'content': reply})

        except Exception as e:
            print(f"\nERROR: {str(e)}\n")

if __name__ == '__main__':
    main()