#!/usr/bin/env python3
"""
╔══════════════════════════════════════════════════════════════╗
║  AGENT — Marziel AI Autonomous Reasoning Engine             ║
║                                                             ║
║  OpenClaw-like capabilities:                                ║
║    • Tool chaining  — picks the right tools, chains results ║
║    • Multi-step     — breaks complex queries into subtasks  ║
║    • Research mode  — autonomous web research               ║
║    • Code execution — generates and runs sandboxed code     ║
║    • Privacy-first  — all actions through Sentinel          ║
╚══════════════════════════════════════════════════════════════╝
"""

import os
import re
import json
import time
import logging
from pathlib import Path
from typing import Optional

logger = logging.getLogger("agent")


# ══════════════════════════════════════════════════════════════
# TOOL REGISTRY — All available tools the agent can use
# ══════════════════════════════════════════════════════════════

class Tool:
    """A tool the agent can invoke."""
    __slots__ = ('name', 'description', 'func', 'parameters', 'category')

    def __init__(self, name: str, description: str, func, parameters: list,
                 category: str = "general"):
        self.name = name
        self.description = description
        self.func = func
        self.parameters = parameters
        self.category = category

    def to_schema(self) -> dict:
        return {
            "name": self.name,
            "description": self.description,
            "parameters": self.parameters,
            "category": self.category,
        }


class ToolRegistry:
    """Registry of all tools available to the agent."""

    def __init__(self):
        self.tools: dict[str, Tool] = {}
        self._register_defaults()

    def register(self, tool: Tool):
        self.tools[tool.name] = tool

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

    def list_schemas(self) -> list:
        return [t.to_schema() for t in self.tools.values()]

    def list_by_category(self) -> dict:
        cats = {}
        for t in self.tools.values():
            cats.setdefault(t.category, []).append(t.to_schema())
        return cats

    def _register_defaults(self):
        """Register all built-in tools."""
        from neurons import (
            synapse_web_search, synapse_fetch_url, synapse_github_repo,
            synapse_weather, synapse_maritime_api, synapse_ml_predict,
            synapse_crawl, synapse_code_exec, SynapsePool,
        )

        pool = SynapsePool(max_workers=4, default_timeout=15)

        def _fire_synapse(func, *args, **kwargs):
            """Fire a synapse worker (spawn → execute → die)."""
            result = pool.fire(func, *args, worker_id=func.__name__, **kwargs)
            if result.success:
                return result.data
            return {"error": result.error}

        # ── Web Intelligence ──
        self.register(Tool(
            name="web_search",
            description="Search the web for information on any topic. Returns results from DuckDuckGo.",
            func=lambda query: _fire_synapse(synapse_web_search, query),
            parameters=[{"name": "query", "type": "string", "required": True}],
            category="intelligence",
        ))

        self.register(Tool(
            name="fetch_url",
            description="Fetch and extract text content from a URL. Strips HTML, returns readable text.",
            func=lambda url: _fire_synapse(synapse_fetch_url, url),
            parameters=[{"name": "url", "type": "string", "required": True}],
            category="intelligence",
        ))

        self.register(Tool(
            name="analyze_github",
            description="Analyze a GitHub repository. Returns stars, forks, language, README preview.",
            func=lambda owner, repo: _fire_synapse(synapse_github_repo, owner, repo),
            parameters=[
                {"name": "owner", "type": "string", "required": True},
                {"name": "repo", "type": "string", "required": True},
            ],
            category="intelligence",
        ))

        self.register(Tool(
            name="deep_crawl",
            description="Crawl a website, following links to gather content from multiple pages.",
            func=lambda url, max_pages=3: _fire_synapse(synapse_crawl, url, max_pages),
            parameters=[
                {"name": "url", "type": "string", "required": True},
                {"name": "max_pages", "type": "int", "default": 3},
            ],
            category="intelligence",
        ))

        # ── Data ──
        self.register(Tool(
            name="weather",
            description="Get 7-day weather and marine forecast for a location.",
            func=lambda lat, lon, name="Unknown": _fire_synapse(synapse_weather, lat, lon, name),
            parameters=[
                {"name": "lat", "type": "float", "required": True},
                {"name": "lon", "type": "float", "required": True},
                {"name": "name", "type": "string", "default": "Unknown"},
            ],
            category="data",
        ))

        self.register(Tool(
            name="maritime_api",
            description="Query the maritime platform API for vessels, ports, voyages, risk zones.",
            func=lambda endpoint: _fire_synapse(synapse_maritime_api, endpoint),
            parameters=[{"name": "endpoint", "type": "string", "required": True}],
            category="data",
        ))

        self.register(Tool(
            name="ml_predict",
            description="Run a machine learning model prediction.",
            func=lambda model, data: _fire_synapse(synapse_ml_predict, model, data),
            parameters=[
                {"name": "model", "type": "string", "required": True},
                {"name": "data", "type": "object", "required": True},
            ],
            category="data",
        ))

        # ── Code ──
        self.register(Tool(
            name="execute_code",
            description="Execute Python code in a sandboxed environment. No network access. 256MB memory limit.",
            func=lambda code: _fire_synapse(synapse_code_exec, code, "python", 5),
            parameters=[{"name": "code", "type": "string", "required": True}],
            category="code",
        ))


# ══════════════════════════════════════════════════════════════
# AGENT PLANNER — Decides which tools to use
# ══════════════════════════════════════════════════════════════

class AgentPlanner:
    """
    Analyzes a query and decides which tools to use.
    No LLM needed for planning — pattern matching + heuristics.
    """

    # Patterns that trigger specific tools
    PATTERNS = {
        "web_search": [
            r"search\s+(for|about|the web)",
            r"(what|who|when|where|how|why)\s+(is|are|was|were|do|does|did)",
            r"(latest|current|recent|news|today)",
            r"(find|look\s+up|google|research)\s+",
            r"(compare|difference\s+between)",
            r"(best\s+(way|practice|tool)|recommend)",
            r"(explain|tell\s+me\s+about)",
        ],
        "fetch_url": [
            r"https?://\S+",
        ],
        "analyze_github": [
            r"github\.com/([\w.-]+)/([\w.-]+)",
            r"(analyze|check|review)\s+(this\s+)?repo",
        ],
        "deep_crawl": [
            r"(crawl|scrape|scan)\s+",
            r"(all\s+pages|entire\s+site|full\s+website)",
        ],
        "weather": [
            r"(weather|forecast|wind|wave|storm|temperature|rain)",
        ],
        "maritime_api": [
            r"(vessel|ship|fleet|port|voyage|route|bunker|risk\s+zone)",
        ],
        "ml_predict": [
            r"(predict|forecast|classify|detect\s+anomal)",
        ],
        "execute_code": [
            r"(run|execute|calculate|compute)\s+(this\s+)?(code|script|program)",
            r"```python",
            r"(solve|calculate)\s+",
        ],
    }

    # Known locations for weather
    LOCATIONS = {
        "rotterdam": (51.92, 4.48), "singapore": (1.29, 103.85),
        "suez": (29.97, 32.53), "gibraltar": (36.14, -5.35),
        "istanbul": (41.02, 29.0), "london": (51.51, -0.13),
        "dubai": (25.27, 55.3), "hamburg": (53.55, 9.99),
        "piraeus": (37.94, 23.65), "shanghai": (31.23, 121.47),
        "houston": (29.76, -95.36), "lisbon": (38.72, -9.14),
        "setubal": (38.52, -8.9), "barcelona": (41.38, 2.17),
        "new york": (40.71, -74.01), "tokyo": (35.68, 139.69),
        "busan": (35.18, 129.07), "los angeles": (33.94, -118.41),
    }

    def plan(self, query: str) -> list:
        """
        Analyze query and return a list of tool calls to execute.
        Returns: [{"tool": "name", "args": {...}}, ...]
        """
        q = query.lower()
        plan = []

        # Check URL patterns first (highest priority)
        urls = re.findall(r'https?://[^\s<>"\')\]]+', query)
        for url in urls[:2]:
            gh_match = re.search(r'github\.com/([\w.-]+)/([\w.-]+)', url)
            if gh_match:
                plan.append({
                    "tool": "analyze_github",
                    "args": {"owner": gh_match.group(1), "repo": gh_match.group(2).rstrip('/')},
                    "reason": f"GitHub repo detected: {gh_match.group(1)}/{gh_match.group(2)}",
                })
            else:
                plan.append({
                    "tool": "fetch_url",
                    "args": {"url": url},
                    "reason": f"URL detected: {url[:60]}",
                })

        # Deep crawl
        if any(re.search(p, q) for p in self.PATTERNS["deep_crawl"]):
            for url in urls[:1]:
                plan.append({
                    "tool": "deep_crawl",
                    "args": {"url": url},
                    "reason": "Deep crawl requested",
                })

        # Weather
        if any(re.search(p, q) for p in self.PATTERNS["weather"]):
            lat, lon, loc = 38.5, -8.9, "Setúbal"  # Default
            for name, (la, lo) in self.LOCATIONS.items():
                if name in q:
                    lat, lon, loc = la, lo, name.title()
                    break
            plan.append({
                "tool": "weather",
                "args": {"lat": lat, "lon": lon, "name": loc},
                "reason": f"Weather for {loc}",
            })

        # Maritime API
        maritime_endpoints = {
            "vessel": "/api/v1/vessel-specs?limit=10",
            "ship": "/api/v1/vessel-specs?limit=10",
            "fleet": "/api/v1/vessel-specs?limit=10",
            "port": "/api/v1/ports?limit=15",
            "voyage": "/api/v1/voyages?limit=10",
            "risk": "/api/v1/risk-zones/active",
            "market": "/api/v1/market-data/prices?limit=20",
            "bunker": "/api/v1/market-data/prices?limit=20",
        }
        for keyword, endpoint in maritime_endpoints.items():
            if keyword in q and not any(t["tool"] == "maritime_api" for t in plan):
                plan.append({
                    "tool": "maritime_api",
                    "args": {"endpoint": endpoint},
                    "reason": f"Maritime data: {keyword}",
                })

        # Code execution
        code_match = re.search(r'```python\n(.*?)```', query, re.DOTALL)
        if code_match:
            plan.append({
                "tool": "execute_code",
                "args": {"code": code_match.group(1)},
                "reason": "Code block detected",
            })

        # ML prediction
        if any(re.search(p, q) for p in self.PATTERNS["ml_predict"]):
            plan.append({
                "tool": "ml_predict",
                "args": {"model": "ensemble_risk", "data": {}},
                "reason": "ML prediction requested",
            })

        # Web search (fallback if no other tools matched)
        if not plan and any(re.search(p, q) for p in self.PATTERNS["web_search"]):
            plan.append({
                "tool": "web_search",
                "args": {"query": query[:120]},
                "reason": "General knowledge query",
            })

        return plan


# ══════════════════════════════════════════════════════════════
# AGENT — The autonomous reasoning engine
# ══════════════════════════════════════════════════════════════

class Agent:
    """
    Marziel Agent — autonomous tool-chaining reasoning engine.

    Flow:
      1. Receive query
      2. Plan (decide tools to use)
      3. Execute tools in parallel via synapses
      4. Synthesize results into context
      5. Return enriched context for LLM
    """

    def __init__(self, tool_registry: ToolRegistry = None):
        self.tools = tool_registry or ToolRegistry()
        self.planner = AgentPlanner()
        self.stats = {
            "total_queries": 0,
            "total_tool_calls": 0,
            "total_plans": 0,
        }

    def process(self, query: str) -> dict:
        """
        Process a query — plan, execute tools, return enriched context.

        Returns:
            {
                "plan": [...],
                "results": [...],
                "context": "synthesized context string",
                "duration_ms": float,
                "tools_used": int,
            }
        """
        t0 = time.monotonic()
        self.stats["total_queries"] += 1

        # 1. Plan — decide which tools to use
        plan = self.planner.plan(query)
        self.stats["total_plans"] += 1

        if not plan:
            return {
                "plan": [],
                "results": [],
                "context": "",
                "duration_ms": 0,
                "tools_used": 0,
            }

        logger.info(f"🤖 Agent plan: {[p['tool'] for p in plan]}")

        # 2. Execute tools (parallel where possible)
        results = []
        for step in plan:
            tool = self.tools.get(step["tool"])
            if not tool:
                results.append({"tool": step["tool"], "error": "Unknown tool"})
                continue

            self.stats["total_tool_calls"] += 1
            try:
                args = step["args"]
                if isinstance(args, dict):
                    result = tool.func(**args)
                else:
                    result = tool.func(args)
                results.append({
                    "tool": step["tool"],
                    "reason": step.get("reason", ""),
                    "data": result,
                })
            except Exception as e:
                results.append({
                    "tool": step["tool"],
                    "error": str(e),
                })

        # 3. Synthesize context
        context = self._synthesize(results)

        duration = (time.monotonic() - t0) * 1000
        logger.info(f"🤖 Agent completed in {duration:.0f}ms ({len(results)} tools)")

        return {
            "plan": plan,
            "results": results,
            "context": context,
            "duration_ms": round(duration, 1),
            "tools_used": len(results),
        }

    def _synthesize(self, results: list) -> str:
        """Synthesize tool results into a context string for the LLM."""
        parts = []

        for r in results:
            if "error" in r and r["error"]:
                continue

            data = r.get("data", {})
            if not data or not isinstance(data, dict):
                continue

            dtype = data.get("type", "")

            if dtype == "web_search":
                ctx = f"WEB SEARCH RESULTS for '{data.get('query', '')}':\n"
                if data.get("abstract"):
                    ctx += f"  Summary: {data['abstract']} (source: {data.get('source', '')})\n"
                for item in data.get("results", [])[:5]:
                    ctx += f"  - {item.get('text', '')} ({item.get('url', '')})\n"
                parts.append(ctx)

            elif dtype == "url_analysis":
                parts.append(f"URL CONTENT ({data.get('url', '')}):\n{data.get('preview', '')[:2000]}")

            elif dtype == "github_analysis":
                d = data
                ctx = (f"GITHUB REPO ({d.get('repo', '')}):\n"
                       f"  ⭐ Stars: {d.get('stars', 0)} | 🍴 Forks: {d.get('forks', 0)}\n"
                       f"  Language: {d.get('language', 'N/A')} | License: {d.get('license', 'N/A')}\n"
                       f"  Issues: {d.get('open_issues', 0)}\n"
                       f"  Description: {d.get('description', 'N/A')}\n")
                if d.get("readme_preview"):
                    ctx += f"  README:\n{d['readme_preview'][:500]}\n"
                parts.append(ctx)

            elif dtype == "weather":
                ctx = f"WEATHER FORECAST for {data.get('location', '')} (7 days):\n"
                for day in data.get("forecast", []):
                    ctx += (f"  {day.get('date', '?')}: {day.get('temp_min', '?')}-{day.get('temp_max', '?')}°C, "
                            f"Wind {day.get('wind_max', '?')} km/h, Rain {day.get('rain_mm', '?')} mm, "
                            f"Waves {day.get('wave_max', '?')}m\n")
                parts.append(ctx)

            elif dtype == "maritime_api":
                ctx = f"MARITIME DATA ({data.get('endpoint', '')}) — {data.get('count', 0)} items:\n"
                for item in data.get("data", [])[:10]:
                    ctx += f"  - {json.dumps(item)[:200]}\n"
                parts.append(ctx)

            elif dtype == "crawl":
                ctx = f"WEB CRAWL ({data.get('start_url', '')}) — {data.get('pages_crawled', 0)} pages:\n"
                for page in data.get("pages", [])[:3]:
                    ctx += f"\n  PAGE: {page.get('url', '')}\n  {page.get('preview', '')[:500]}\n"
                parts.append(ctx)

            elif dtype == "code_exec":
                ctx = "CODE EXECUTION RESULT:\n"
                if data.get("stdout"):
                    ctx += f"  Output:\n{data['stdout'][:2000]}\n"
                if data.get("stderr"):
                    ctx += f"  Errors:\n{data['stderr'][:500]}\n"
                parts.append(ctx)

            elif dtype == "ml_predict":
                parts.append(f"ML PREDICTION ({data.get('model', '')}):\n  {json.dumps(data.get('prediction', {}))[:500]}")

        return "\n".join(parts)

    def get_stats(self) -> dict:
        return dict(self.stats)


# ══════════════════════════════════════════════════════════════
# CONVENIENCE: Global agent instance
# ══════════════════════════════════════════════════════════════

_agent = None

def get_agent() -> Agent:
    """Get or create the global Agent instance."""
    global _agent
    if _agent is None:
        _agent = Agent()
    return _agent
