#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
LLM-powered task decomposition for PULSE dispatcher.
Extracted from pulse_dispatcher.py per Law 7.

Provider-agnostic LLM abstraction with SSRF protection.
Dependencies: Python stdlib only (Law 4)
"""
import json
import os
import re
import urllib.request
import urllib.error
from datetime import datetime, timezone
from pathlib import Path
from pulse.core.brain_utils import now_iso as _now_iso


def _ts():
    return int(datetime.now(timezone.utc).timestamp())


# G11: Allowed endpoint hosts for SSRF prevention
_ALLOWED_HOSTS = frozenset({
    "api.anthropic.com", "api.openai.com",
    "generativelanguage.googleapis.com", "api.x.ai",
    "localhost", "127.0.0.1",
})

# G6: Valid dispatcher providers
VALID_PROVIDERS = frozenset({
    "none", "anthropic", "google", "openai", "ollama", "grok",
})

# G2: Valid domains and complexities — reject unknown values
_VALID_DOMAINS = frozenset({
    "security", "architecture", "testing", "features",
    "documentation", "refactoring", "infrastructure",
})
_VALID_COMPLEXITY = frozenset({"low", "medium", "high"})
_MAX_FIELD_LEN = 500  # Max chars for title/description

# Smart Routing: complexity → model tier
_TIER_MAP = {
    "low": "fast",       # Haiku, Flash, Gemini Flash — cheap/fast
    "medium": "standard", # Sonnet, Gemini Pro — balanced
    "high": "premium",    # Opus, GPT-4, Gemini Ultra — expensive/smart
}


def _validate_endpoint(url: str) -> bool:
    """G11: Validate URL host against whitelist to prevent SSRF."""
    try:
        from urllib.parse import urlparse
        parsed = urlparse(url)
        host = parsed.hostname or ""
        return host in _ALLOWED_HOSTS
    except Exception:
        return False


def _sanitize_text(text: str) -> str:
    """G2: Sanitize text fields — strip control chars, limit length."""
    if not isinstance(text, str):
        return ""
    # Strip control characters except newlines
    cleaned = re.sub(r'[\x00-\x09\x0b-\x1f\x7f]', '', text)
    return cleaned[:_MAX_FIELD_LEN]


def _normalize_tasks(raw_tasks: list) -> list[dict]:
    """Normalize LLM output into valid task schema.
    G2: Sanitizes all text fields. G5: Validates domains/complexity."""
    ts = _ts()
    normalized = []
    for i, t in enumerate(raw_tasks):
        if not isinstance(t, dict):
            continue
        # Convert index-based depends_on to task_id-based
        deps = []
        for dep in t.get("depends_on", []):
            if isinstance(dep, int) and 0 <= dep < len(raw_tasks):
                deps.append(f"t_{ts}_{dep + 1:03d}")

        # G2: Sanitize text fields
        title = _sanitize_text(t.get("title", f"Task {i + 1}"))
        description = _sanitize_text(t.get("description", ""))

        # G5: Validate domain and complexity against whitelist
        raw_domain = t.get("domain", "features")
        domain = raw_domain if raw_domain in _VALID_DOMAINS else "features"
        raw_complexity = t.get("complexity", "medium")
        complexity = raw_complexity if raw_complexity in _VALID_COMPLEXITY else "medium"

        # G2: Sanitize capabilities list (strings only, no special chars)
        raw_caps = t.get("required_capabilities", [])
        capabilities = [
            re.sub(r'[^a-zA-Z0-9_\-]', '', str(c))[:50]
            for c in raw_caps if isinstance(c, str)
        ][:10]  # Cap at 10 capabilities

        task = {
            "task_id": f"t_{ts}_{i + 1:03d}",
            "title": title,
            "description": description,
            "domain": domain,
            "complexity": complexity,
            "recommended_tier": _TIER_MAP.get(complexity, "standard"),
            "required_capabilities": capabilities,
            "depends_on": deps,
            "status": "blocked" if deps else "open",
            "claimed_by": None,
            "claimed_at": None,
            "files_touched": [],
            "outcome": None,
        }
        normalized.append(task)
    return normalized


def _llm_call(provider: str, model: str, endpoint: str, system_prompt: str, user_prompt: str, verbose=False) -> list[dict]:
    """Provider-agnostic LLM call. Returns parsed task list or empty list."""
    if provider == "none" or not provider:
        return []

    # Build OpenAI-compatible request (works for most providers)
    if provider == "ollama":
        url = (endpoint or "http://localhost:11434") + "/api/chat"
        body = {
            "model": model or "llama3",
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            "stream": False,
            "format": "json",
        }
    else:
        endpoints = {
            "anthropic": "https://api.anthropic.com/v1/messages",
            "openai": "https://api.openai.com/v1/chat/completions",
            "google": "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
            "grok": "https://api.x.ai/v1/chat/completions",
        }
        url = endpoint or endpoints.get(provider, "")
        if not url:
            return []

        if provider == "anthropic":
            body = {
                "model": model or "claude-haiku-4-5-20251001",
                "max_tokens": 2048,
                "system": system_prompt,
                "messages": [{"role": "user", "content": user_prompt}],
            }
        else:
            body = {
                "model": model or "gpt-4o-mini",
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                "response_format": {"type": "json_object"},
            }

    # Build headers
    headers = {"Content-Type": "application/json"}
    if provider == "anthropic":
        api_key = os.environ.get("ANTHROPIC_API_KEY", "")
        headers["x-api-key"] = api_key
        headers["anthropic-version"] = "2023-06-01"
    elif provider == "openai":
        api_key = os.environ.get("OPENAI_API_KEY", "")
        headers["Authorization"] = f"Bearer {api_key}"
    elif provider == "google":
        api_key = os.environ.get("GOOGLE_API_KEY", "")
        headers["Authorization"] = f"Bearer {api_key}"
    elif provider == "grok":
        api_key = os.environ.get("XAI_API_KEY", "")
        headers["Authorization"] = f"Bearer {api_key}"

    # G11: SSRF prevention — validate endpoint URL
    if not _validate_endpoint(url):
        if verbose:
            print(f"[DECOMPOSER] SECURITY: Blocked request to unauthorized host in URL")
        return []

    try:
        data = json.dumps(body).encode("utf-8")
        req = urllib.request.Request(url, data=data, headers=headers, method="POST")
        with urllib.request.urlopen(req, timeout=30) as resp:
            result = json.loads(resp.read().decode("utf-8"))

        # Extract content based on provider
        if provider == "anthropic":
            text = result.get("content", [{}])[0].get("text", "")
        elif provider == "ollama":
            text = result.get("message", {}).get("content", "")
        else:
            text = result.get("choices", [{}])[0].get("message", {}).get("content", "")

        # Parse JSON from response
        parsed = json.loads(text)
        if isinstance(parsed, list):
            return _normalize_tasks(parsed)
        if isinstance(parsed, dict) and "tasks" in parsed:
            return _normalize_tasks(parsed["tasks"])
        return []

    except (urllib.error.URLError, json.JSONDecodeError, KeyError, IndexError) as e:
        if verbose:
            print(f"[DECOMPOSER] LLM call failed ({provider}): {e}")
        return []


def decompose_prompt(prompt: str, provider: str = "none", model: str = "", endpoint: str = "", verbose: bool = False) -> list[dict]:
    """Use configured LLM to decompose a prompt into tasks.
    Returns list of task dicts. Falls back to empty list on failure.

    Args:
        prompt: User's task description
        provider: One of: none, anthropic, google, openai, ollama, grok
        model: Model name (provider-specific)
        endpoint: Optional custom endpoint URL
        verbose: Enable debug logging

    Returns:
        List of normalized task dicts with schema-compliant fields
    """
    # G6: Validate provider against whitelist
    validated_provider = provider if provider in VALID_PROVIDERS else "none"
    if provider != validated_provider and verbose:
        print(f"[DECOMPOSER] WARNING: Invalid provider '{provider}', falling back to 'none'")

    if validated_provider == "none":
        return []

    system_msg = (
        "You are a task decomposer for PULSE, a multi-agent AI system. "
        "Break the user's prompt into independent, parallelizable tasks. "
        "Return ONLY a JSON array of objects with fields: "
        "title, description, domain, complexity (low/medium/high), "
        "required_capabilities (array), depends_on (array of indices, 0-based). "
        "Domains: security, architecture, testing, features, documentation, refactoring."
    )

    return _llm_call(validated_provider, model, endpoint, system_msg, prompt, verbose)
