"""
Doit Agent - Tool Registry & Implementations
All tools that Doit can execute. AI can ONLY invoke registered tools.
"""
from __future__ import annotations

import asyncio
import logging
import os
import platform
import shutil
import subprocess
import sys
import time
import zipfile
import tarfile
from pathlib import Path
from typing import Any, Callable, Optional

import aiofiles
import aiohttp
import psutil

from core.config import (
    MAX_CPU_PERCENT, MAX_DISK_WRITE_MB, HTTP_TIMEOUT, BLOCKED_PATHS
)
from security.security import InjectionGuard

logger = logging.getLogger("doit.tools")

# Tool definition
class ToolDefinition:
    def __init__(self, name: str, description: str, fn: Callable, dangerous: bool = False):
        self.name = name
        self.description = description
        self.fn = fn
        self.dangerous = dangerous


# ── File System Tools ────────────────────────────────────────────────────────

async def fs_list(path: str = ".", **_) -> dict:
    p = Path(path).expanduser().resolve()
    safe, reason = InjectionGuard.check_path(str(p))
    if not safe:
        return {"error": reason}
    if not p.exists():
        return {"error": f"Path not found: {p}"}
    if p.is_file():
        return {"type": "file", "path": str(p), "size": p.stat().st_size}
    entries = []
    for item in sorted(p.iterdir()):
        try:
            stat = item.stat()
            entries.append({
                "name": item.name,
                "type": "dir" if item.is_dir() else "file",
                "size": stat.st_size if item.is_file() else None,
            })
        except PermissionError:
            entries.append({"name": item.name, "type": "unknown", "error": "permission denied"})
    return {"path": str(p), "entries": entries, "count": len(entries)}


async def fs_read(path: str, **_) -> dict:
    p = Path(path).expanduser().resolve()
    safe, reason = InjectionGuard.check_path(str(p))
    if not safe:
        return {"error": reason}
    if not p.exists():
        return {"error": f"File not found: {p}"}
    if p.stat().st_size > 5 * 1024 * 1024:  # 5MB limit
        return {"error": "File too large to read (>5MB)"}
    try:
        async with aiofiles.open(p, "r", encoding="utf-8", errors="replace") as f:
            content = await f.read()
        return {"path": str(p), "content": content, "size": len(content)}
    except Exception as e:
        return {"error": str(e)}


async def fs_write(path: str, content: str, **_) -> dict:
    p = Path(path).expanduser().resolve()
    safe, reason = InjectionGuard.check_path(str(p))
    if not safe:
        return {"error": reason}
    size_bytes = len(content.encode("utf-8"))
    if size_bytes > MAX_DISK_WRITE_MB * 1024 * 1024:
        return {"error": f"Content too large (>{MAX_DISK_WRITE_MB}MB)"}
    try:
        p.parent.mkdir(parents=True, exist_ok=True)
        async with aiofiles.open(p, "w", encoding="utf-8") as f:
            await f.write(content)
        return {"path": str(p), "bytes_written": size_bytes, "success": True}
    except Exception as e:
        return {"error": str(e)}


async def fs_delete(path: str, **_) -> dict:
    p = Path(path).expanduser().resolve()
    safe, reason = InjectionGuard.check_path(str(p))
    if not safe:
        return {"error": reason}
    if not p.exists():
        return {"error": f"Path not found: {p}"}
    try:
        if p.is_dir():
            shutil.rmtree(p)
        else:
            p.unlink()
        return {"deleted": str(p), "success": True}
    except Exception as e:
        return {"error": str(e)}


async def fs_move(src: str, dst: str, **_) -> dict:
    s = Path(src).expanduser().resolve()
    d = Path(dst).expanduser().resolve()
    for p in (s, d):
        safe, reason = InjectionGuard.check_path(str(p))
        if not safe:
            return {"error": reason}
    if not s.exists():
        return {"error": f"Source not found: {s}"}
    try:
        d.parent.mkdir(parents=True, exist_ok=True)
        shutil.move(str(s), str(d))
        return {"src": str(s), "dst": str(d), "success": True}
    except Exception as e:
        return {"error": str(e)}


async def fs_copy(src: str, dst: str, **_) -> dict:
    s = Path(src).expanduser().resolve()
    d = Path(dst).expanduser().resolve()
    for p in (s, d):
        safe, reason = InjectionGuard.check_path(str(p))
        if not safe:
            return {"error": reason}
    if not s.exists():
        return {"error": f"Source not found: {s}"}
    try:
        d.parent.mkdir(parents=True, exist_ok=True)
        if s.is_dir():
            shutil.copytree(str(s), str(d))
        else:
            shutil.copy2(str(s), str(d))
        return {"src": str(s), "dst": str(d), "success": True}
    except Exception as e:
        return {"error": str(e)}


async def fs_search(path: str, pattern: str, **_) -> dict:
    p = Path(path).expanduser().resolve()
    safe, reason = InjectionGuard.check_path(str(p))
    if not safe:
        return {"error": reason}
    if not p.exists():
        return {"error": f"Path not found: {p}"}
    results = []
    try:
        for match in p.rglob(pattern):
            results.append({
                "path": str(match),
                "type": "dir" if match.is_dir() else "file",
                "size": match.stat().st_size if match.is_file() else None,
            })
            if len(results) >= 100:  # Limit results
                break
    except Exception as e:
        return {"error": str(e)}
    return {"pattern": pattern, "results": results, "count": len(results)}


async def fs_compress(src: str, dst: str, **_) -> dict:
    s = Path(src).expanduser().resolve()
    d = Path(dst).expanduser().resolve()
    for p in (s, d):
        safe, reason = InjectionGuard.check_path(str(p))
        if not safe:
            return {"error": reason}
    if not s.exists():
        return {"error": f"Source not found: {s}"}
    try:
        d.parent.mkdir(parents=True, exist_ok=True)
        if dst.endswith(".tar.gz") or dst.endswith(".tgz"):
            with tarfile.open(str(d), "w:gz") as tar:
                tar.add(str(s), arcname=s.name)
        else:
            with zipfile.ZipFile(str(d), "w", zipfile.ZIP_DEFLATED) as zf:
                if s.is_dir():
                    for f in s.rglob("*"):
                        zf.write(f, f.relative_to(s.parent))
                else:
                    zf.write(str(s), s.name)
        return {"src": str(s), "dst": str(d), "size": d.stat().st_size, "success": True}
    except Exception as e:
        return {"error": str(e)}


async def fs_extract(src: str, dst: str, **_) -> dict:
    s = Path(src).expanduser().resolve()
    d = Path(dst).expanduser().resolve()
    for p in (s, d):
        safe, reason = InjectionGuard.check_path(str(p))
        if not safe:
            return {"error": reason}
    if not s.exists():
        return {"error": f"Archive not found: {s}"}
    try:
        d.mkdir(parents=True, exist_ok=True)
        if tarfile.is_tarfile(str(s)):
            with tarfile.open(str(s), "r:*") as tar:
                tar.extractall(str(d))
        elif zipfile.is_zipfile(str(s)):
            with zipfile.ZipFile(str(s), "r") as zf:
                zf.extractall(str(d))
        else:
            return {"error": "Unsupported archive format"}
        return {"src": str(s), "dst": str(d), "success": True}
    except Exception as e:
        return {"error": str(e)}


ORGANIZE_RULES = {
    "images": [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".svg", ".webp", ".heic"],
    "documents": [".pdf", ".doc", ".docx", ".xls", ".xlsx", ".ppt", ".pptx", ".txt", ".md"],
    "videos": [".mp4", ".avi", ".mkv", ".mov", ".wmv", ".flv", ".webm"],
    "audio": [".mp3", ".wav", ".flac", ".aac", ".ogg", ".m4a"],
    "archives": [".zip", ".tar", ".gz", ".7z", ".rar", ".bz2"],
    "code": [".py", ".js", ".ts", ".html", ".css", ".java", ".c", ".cpp", ".go", ".rs"],
}


async def fs_organize(path: str, **_) -> dict:
    p = Path(path).expanduser().resolve()
    safe, reason = InjectionGuard.check_path(str(p))
    if not safe:
        return {"error": reason}
    if not p.is_dir():
        return {"error": f"Not a directory: {p}"}
    moved = []
    for item in p.iterdir():
        if item.is_file():
            ext = item.suffix.lower()
            for folder, exts in ORGANIZE_RULES.items():
                if ext in exts:
                    target_dir = p / folder
                    target_dir.mkdir(exist_ok=True)
                    shutil.move(str(item), str(target_dir / item.name))
                    moved.append({"file": item.name, "to": folder})
                    break
    return {"path": str(p), "moved": moved, "count": len(moved), "success": True}


async def fs_backup(src: str, dst: str, **_) -> dict:
    from datetime import datetime as _dt
    s = Path(src).expanduser().resolve()
    d = Path(dst).expanduser().resolve()
    for p, name in ((s, "src"), (d, "dst")):
        safe, reason = InjectionGuard.check_path(str(p))
        if not safe:
            return {"error": f"Blocked {name}: {reason}"}
    if not s.exists():
        return {"error": f"Source not found: {s}"}
    ts = _dt.now().strftime("%Y%m%d_%H%M%S")
    backup_name = f"{s.name}_backup_{ts}"
    backup_path = d / backup_name
    result = await fs_copy(str(s), str(backup_path))
    if "error" not in result:
        result["backup_path"] = str(backup_path)
        result["timestamp"] = ts
    return result


async def fs_clean(path: str, days: int = None, size_mb: float = None,
                   type: str = None, **_) -> dict:
    import time as t
    p = Path(path).expanduser().resolve()
    safe, reason = InjectionGuard.check_path(str(p))
    if not safe:
        return {"error": reason}
    if not p.is_dir():
        return {"error": f"Not a directory: {p}"}
    deleted = []
    now = t.time()
    for item in p.iterdir():
        if not item.is_file():
            continue
        try:
            stat = item.stat()
            should_delete = False
            if days and (now - stat.st_mtime) > days * 86400:
                should_delete = True
            if size_mb and stat.st_size > size_mb * 1024 * 1024:
                should_delete = True
            if type and item.suffix.lower() == f".{type.lstrip('.')}":
                should_delete = True
            if should_delete:
                item.unlink()
                deleted.append(item.name)
        except Exception:
            pass
    return {"deleted": deleted, "count": len(deleted), "success": True}


# ── System Tools ─────────────────────────────────────────────────────────────

async def sys_info(**_) -> dict:
    return {
        "os": platform.system(),
        "os_version": platform.version(),
        "hostname": platform.node(),
        "python": platform.python_version(),
        "cpu_count": psutil.cpu_count(),
        "cpu_freq_mhz": psutil.cpu_freq().current if psutil.cpu_freq() else None,
        "ram_total_gb": round(psutil.virtual_memory().total / 1e9, 2),
        "disk_total_gb": round(psutil.disk_usage("/").total / 1e9, 2),
        "uptime_hours": round((time.time() - psutil.boot_time()) / 3600, 1),
        "architecture": platform.machine(),
    }


async def sys_processes(**_) -> dict:
    procs = []
    for p in psutil.process_iter(["pid", "name", "status", "cpu_percent", "memory_info"]):
        try:
            procs.append({
                "pid": p.info["pid"],
                "name": p.info["name"],
                "status": p.info["status"],
                "cpu_pct": p.info["cpu_percent"],
                "mem_mb": round(p.info["memory_info"].rss / 1e6, 1) if p.info["memory_info"] else None,
            })
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
    procs.sort(key=lambda x: x.get("mem_mb") or 0, reverse=True)
    return {"processes": procs[:50], "total": len(procs)}


async def sys_kill(pid: int, **_) -> dict:
    # Protect system-critical processes
    try:
        proc = psutil.Process(pid)
        name = proc.name().lower()
        critical = ["systemd", "init", "kernel", "kthreadd", "launchd", "svchost"]
        if any(c in name for c in critical):
            return {"error": f"Refusing to kill critical process: {name}"}
        proc.terminate()
        try:
            proc.wait(timeout=5)
        except psutil.TimeoutExpired:
            proc.kill()
        return {"pid": pid, "success": True}
    except psutil.NoSuchProcess:
        return {"error": f"Process {pid} not found"}
    except psutil.AccessDenied:
        return {"error": f"Access denied to process {pid}"}


async def sys_monitor(**_) -> dict:
    vm = psutil.virtual_memory()
    disk = psutil.disk_usage("/")
    return {
        "cpu_percent": psutil.cpu_percent(interval=1),
        "ram_used_gb": round(vm.used / 1e9, 2),
        "ram_total_gb": round(vm.total / 1e9, 2),
        "ram_percent": vm.percent,
        "disk_used_gb": round(disk.used / 1e9, 2),
        "disk_total_gb": round(disk.total / 1e9, 2),
        "disk_percent": disk.percent,
        "net_sent_mb": round(psutil.net_io_counters().bytes_sent / 1e6, 1),
        "net_recv_mb": round(psutil.net_io_counters().bytes_recv / 1e6, 1),
    }


async def sys_shutdown(delay_s: int = 0, **_) -> dict:
    if delay_s > 0:
        await asyncio.sleep(delay_s)
    try:
        if sys.platform == "win32":
            result = subprocess.run(["shutdown", "/s", "/t", "0"],
                                    capture_output=True, text=True)
            if result.returncode != 0:
                return {"error": result.stderr}
        elif sys.platform == "darwin":
            result = subprocess.run(["osascript", "-e",
                                     'tell app "System Events" to shut down'],
                                    capture_output=True, text=True)
            if result.returncode != 0:
                # Fallback: try shutdown with sudo (may prompt)
                result = subprocess.run(["sudo", "shutdown", "-h", "now"],
                                        capture_output=True, text=True)
                if result.returncode != 0:
                    return {"error": "Shutdown failed. May need admin rights."}
        else:
            # Linux: try systemctl first (works in user session on systemd)
            result = subprocess.run(["systemctl", "poweroff"],
                                    capture_output=True, text=True)
            if result.returncode != 0:
                # Fallback: shutdown command (needs sudo on most systems)
                result = subprocess.run(["sudo", "shutdown", "-h", "now"],
                                        capture_output=True, text=True)
                if result.returncode != 0:
                    return {"error": f"Shutdown failed: {result.stderr}. Try: sudo shutdown -h now"}
        return {"success": True, "action": "shutdown"}
    except Exception as e:
        return {"error": str(e)}


async def sys_restart(delay_s: int = 0, **_) -> dict:
    if delay_s > 0:
        await asyncio.sleep(delay_s)
    try:
        if sys.platform == "win32":
            result = subprocess.run(["shutdown", "/r", "/t", "0"],
                                    capture_output=True, text=True)
            if result.returncode != 0:
                return {"error": result.stderr}
        elif sys.platform == "darwin":
            result = subprocess.run(["osascript", "-e",
                                     'tell app "System Events" to restart'],
                                    capture_output=True, text=True)
            if result.returncode != 0:
                result = subprocess.run(["sudo", "shutdown", "-r", "now"],
                                        capture_output=True, text=True)
                if result.returncode != 0:
                    return {"error": "Restart failed. May need admin rights."}
        else:
            result = subprocess.run(["systemctl", "reboot"],
                                    capture_output=True, text=True)
            if result.returncode != 0:
                result = subprocess.run(["sudo", "shutdown", "-r", "now"],
                                        capture_output=True, text=True)
                if result.returncode != 0:
                    return {"error": f"Restart failed: {result.stderr}"}
        return {"success": True, "action": "restart"}
    except Exception as e:
        return {"error": str(e)}


async def sys_sleep(**_) -> dict:
    try:
        if sys.platform == "win32":
            result = subprocess.run(
                ["rundll32.exe", "powrprof.dll,SetSuspendState", "0,1,0"],
                capture_output=True, text=True
            )
        elif sys.platform == "darwin":
            result = subprocess.run(["pmset", "sleepnow"],
                                    capture_output=True, text=True)
        else:
            result = subprocess.run(["systemctl", "suspend"],
                                    capture_output=True, text=True)
            if result.returncode != 0:
                result = subprocess.run(["sudo", "systemctl", "suspend"],
                                        capture_output=True, text=True)
        return {"success": True, "action": "sleep"}
    except Exception as e:
        return {"error": str(e)}


# ── Network Tools ─────────────────────────────────────────────────────────────

async def net_download(url: str, dst: str, **_) -> dict:
    d = Path(dst).expanduser().resolve()
    safe, reason = InjectionGuard.check_path(str(d))
    if not safe:
        return {"error": reason}
    d.parent.mkdir(parents=True, exist_ok=True)
    max_bytes = MAX_DISK_WRITE_MB * 1024 * 1024
    try:
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=300)) as sess:
            async with sess.get(url) as resp:
                resp.raise_for_status()
                # Check Content-Length header first if available
                content_length = resp.headers.get("Content-Length")
                if content_length and int(content_length) > max_bytes:
                    return {"error": f"File too large (>{MAX_DISK_WRITE_MB}MB according to Content-Length)"}
                bytes_written = 0
                async with aiofiles.open(d, "wb") as f:
                    async for chunk in resp.content.iter_chunked(65536):  # 64KB chunks
                        bytes_written += len(chunk)
                        if bytes_written > max_bytes:
                            # Delete partial file
                            await f.flush()
                        if bytes_written > max_bytes:
                            try:
                                d.unlink()
                            except Exception:
                                pass
                            return {"error": f"File too large (>{MAX_DISK_WRITE_MB}MB), download aborted"}
                        await f.write(chunk)
        return {"url": url, "dst": str(d), "size_bytes": bytes_written, "success": True}
    except aiohttp.ClientResponseError as e:
        return {"error": f"HTTP {e.status}: {e.message}"}
    except Exception as e:
        return {"error": str(e)}


async def net_ping(host: str, **_) -> dict:
    cmd = ["ping", "-c", "4", host] if sys.platform != "win32" else ["ping", "-n", "4", host]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
        return {
            "host": host,
            "success": result.returncode == 0,
            "output": result.stdout[-500:],
        }
    except subprocess.TimeoutExpired:
        return {"error": "Ping timeout"}
    except Exception as e:
        return {"error": str(e)}


async def net_http(method: str, url: str, headers: dict = None, body: Any = None, **_) -> dict:
    method = method.upper()
    if method not in ("GET", "POST", "PUT", "PATCH", "DELETE", "HEAD"):
        return {"error": f"Unsupported method: {method}"}
    try:
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=HTTP_TIMEOUT)) as sess:
            kwargs = {"headers": headers or {}}
            if body:
                kwargs["json"] = body
            async with sess.request(method, url, **kwargs) as resp:
                text = await resp.text()
                return {
                    "status": resp.status,
                    "headers": dict(resp.headers),
                    "body": text[:5000],
                    "success": 200 <= resp.status < 300,
                }
    except Exception as e:
        return {"error": str(e)}


async def net_monitor(url: str, **_) -> dict:
    try:
        start = time.time()
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=10)) as sess:
            async with sess.get(url) as resp:
                elapsed = round((time.time() - start) * 1000, 1)
                return {
                    "url": url,
                    "status": resp.status,
                    "up": resp.status < 400,
                    "latency_ms": elapsed,
                }
    except Exception as e:
        return {"url": url, "up": False, "error": str(e)}


# ── Tool Registry ─────────────────────────────────────────────────────────────

class ToolRegistry:
    """Central registry of all available tools. AI cannot invoke unregistered tools."""

    def __init__(self):
        self._tools: dict[str, ToolDefinition] = {}
        self._register_builtin()

    def _register_builtin(self):
        tools = [
            ("fs_list", "List directory contents", fs_list, False),
            ("fs_read", "Read file contents", fs_read, False),
            ("fs_write", "Write/create file", fs_write, False),
            ("fs_delete", "Delete file or directory", fs_delete, True),
            ("fs_move", "Move/rename file", fs_move, False),
            ("fs_copy", "Copy file", fs_copy, False),
            ("fs_search", "Search for files", fs_search, False),
            ("fs_compress", "Compress files", fs_compress, False),
            ("fs_extract", "Extract archive", fs_extract, False),
            ("fs_organize", "Auto-organize directory", fs_organize, False),
            ("fs_backup", "Backup path", fs_backup, False),
            ("fs_clean", "Delete old files", fs_clean, True),
            ("sys_info", "Get system information", sys_info, False),
            ("sys_processes", "List running processes", sys_processes, False),
            ("sys_kill", "Kill process", sys_kill, True),
            ("sys_monitor", "Get resource usage", sys_monitor, False),
            ("sys_shutdown", "Shutdown system", sys_shutdown, True),
            ("sys_restart", "Restart system", sys_restart, True),
            ("sys_sleep", "Sleep/suspend system", sys_sleep, False),
            ("net_download", "Download URL", net_download, False),
            ("net_ping", "Ping host", net_ping, False),
            ("net_http", "Make HTTP request", net_http, False),
            ("net_monitor", "Check URL uptime", net_monitor, False),
        ]
        for name, desc, fn, dangerous in tools:
            self.register(name, desc, fn, dangerous)

    def register(self, name: str, description: str, fn: Callable, dangerous: bool = False):
        self._tools[name] = ToolDefinition(name, description, fn, dangerous)

    def get(self, name: str) -> Optional[ToolDefinition]:
        return self._tools.get(name)

    def exists(self, name: str) -> bool:
        return name in self._tools

    def all_names(self) -> list[str]:
        return list(self._tools.keys())

    async def execute(self, tool_name: str, args: dict) -> dict:
        """Execute a tool with strict validation."""
        if not self.exists(tool_name):
            return {"error": f"Unknown tool: {tool_name}. Tools cannot be invented."}
        tool = self.get(tool_name)
        try:
            result = await tool.fn(**args)
            return result
        except TypeError as e:
            return {"error": f"Invalid arguments for {tool_name}: {e}"}
        except Exception as e:
            logger.exception("Tool %s raised exception", tool_name)
            return {"error": f"Tool error: {e}"}


# Singleton
_registry: Optional[ToolRegistry] = None


def get_registry() -> ToolRegistry:
    global _registry
    if _registry is None:
        _registry = ToolRegistry()
    return _registry
