"""
Sulphur Prompt Enhancer Integration for NEXUS OS v2
Uses hugging-science/sulphur_prompt_enhancer-Q4_K_M-imatrix.gguf

The Sulphur model rewrites prompts for better model routing and quality.
It runs before ChimeraRouter to improve intent classification accuracy.
"""
from typing import Optional, Dict, Any
from dataclasses import dataclass
import json
import subprocess


@dataclass
class EnhancedPrompt:
    original: str
    enhanced: str
    intent_tags: list[str]           # Extracted intent labels
    complexity_score: float           # 0-1 estimated reasoning depth
    suggested_tier: str               # "local_8gb", "local_16gb", "cloud_api", etc.
    confidence: float                 # Enhancer confidence


class SulphurEnhancer:
    """
    Sulphur Prompt Enhancer wrapper for Ollama.
    Model: hugging-science/sulphur_prompt_enhancer-Q4_K_M-imatrix.gguf
    """

    MODEL_TAG = "hugging-science/sulphur_prompt_enhancer-Q4_K_M-imatrix.gguf"
    DEFAULT_SYSTEM = """You are a prompt enhancement specialist. Your task:
1. Analyze the user's intent, complexity, and required capabilities
2. Rewrite the prompt to be clearer, more specific, and better structured
3. Tag with: [coding], [reasoning], [vision], [creative], [factual], [safety], [fast], [long_context]
4. Estimate complexity (0-1) and suggest target tier (8gb/16gb/24gb/48gb/cloud)
5. Output ONLY valid JSON with keys: enhanced, tags, complexity, tier, confidence"""

    def __init__(
        self,
        model_tag: Optional[str] = None,
        ollama_host: str = "http://localhost:11434",
        temperature: float = 0.3,
    ):
        self.model_tag = model_tag or self.MODEL_TAG
        self.ollama_host = ollama_host
        self.temperature = temperature

    def _call_ollama(self, prompt: str) -> str:
        """Call Sulphur model via Ollama API."""
        import urllib.request
        import urllib.error

        payload = json.dumps({
            "model": self.model_tag,
            "messages": [
                {"role": "system", "content": self.DEFAULT_SYSTEM},
                {"role": "user", "content": f"Enhance this prompt:\n{prompt}"},
            ],
            "stream": False,
            "options": {"temperature": self.temperature},
        }).encode("utf-8")

        req = urllib.request.Request(
            f"{self.ollama_host}/api/chat",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST",
        )

        try:
            with urllib.request.urlopen(req, timeout=120) as resp:
                data = json.loads(resp.read().decode("utf-8"))
                return data.get("message", {}).get("content", "")
        except urllib.error.URLError as e:
            raise RuntimeError(f"Ollama unreachable at {self.ollama_host}: {e}")
        except json.JSONDecodeError:
            raise RuntimeError("Invalid JSON from Ollama API")

    def enhance(self, prompt: str) -> EnhancedPrompt:
        """Enhance a raw user prompt via Sulphur model."""
        raw = self._call_ollama(prompt)

        # Extract JSON from response (may be wrapped in markdown)
        json_match = None
        for pattern in [r'```json\s*(.*?)\s*```', r'```\s*(.*?)\s*```', r'\{.*\}']:
            import re
            m = re.search(pattern, raw, re.DOTALL)
            if m:
                json_match = m.group(1) if m.groups() else m.group(0)
                break

        try:
            data = json.loads(json_match or raw)
        except (json.JSONDecodeError, TypeError):
            # Fallback: parse heuristically
            data = self._heuristic_parse(raw, prompt)

        return EnhancedPrompt(
            original=prompt,
            enhanced=data.get("enhanced", prompt),
            intent_tags=data.get("tags", []),
            complexity_score=float(data.get("complexity", 0.5)),
            suggested_tier=data.get("tier", "local_16gb"),
            confidence=float(data.get("confidence", 0.7)),
        )

    def _heuristic_parse(self, raw: str, original: str) -> Dict[str, Any]:
        """Fallback parser when JSON is malformed."""
        # Simple keyword detection
        tags = []
        if any(kw in raw.lower() for kw in ["code", "program", "function", "debug", "syntax"]):
            tags.append("coding")
        if any(kw in raw.lower() for kw in ["think", "reason", "logic", "step", "analysis"]):
            tags.append("reasoning")
        if any(kw in raw.lower() for kw in ["image", "vision", "see", "photo", "diagram"]):
            tags.append("vision")
        if any(kw in raw.lower() for kw in ["create", "write", "story", "poem", "design"]):
            tags.append("creative")
        if any(kw in raw.lower() for kw in ["fact", "who", "when", "where", "what is"]):
            tags.append("factual")

        complexity = 0.7 if "reasoning" in tags else (0.5 if "coding" in tags else 0.3)
        tier = "cloud_api" if complexity > 0.8 else ("local_24gb" if complexity > 0.6 else "local_16gb")

        return {
            "enhanced": raw.strip() or original,
            "tags": tags or ["general"],
            "complexity": complexity,
            "tier": tier,
            "confidence": 0.5,
        }

    def quick_enhance(self, prompt: str) -> str:
        """Return just the enhanced prompt text."""
        return self.enhance(prompt).enhanced


class MockSulphurEnhancer:
    """Offline mock for testing without Ollama."""

    def enhance(self, prompt: str) -> EnhancedPrompt:
        tags = []
        if any(kw in prompt.lower() for kw in ["code", "python", "function"]):
            tags.append("coding")
        if any(kw in prompt.lower() for kw in ["why", "how", "explain", "reason"]):
            tags.append("reasoning")
        if any(kw in prompt.lower() for kw in ["image", "picture", "vision"]):
            tags.append("vision")
        if any(kw in prompt.lower() for kw in ["write", "create", "story"]):
            tags.append("creative")

        complexity = 0.8 if "reasoning" in tags else (0.6 if "coding" in tags else 0.4)
        tier = "cloud_api" if complexity > 0.8 and "reasoning" in tags else "local_16gb"

        return EnhancedPrompt(
            original=prompt,
            enhanced=f"[ENHANCED] {prompt}\n[Intent: {', '.join(tags or ['general'])}]",
            intent_tags=tags or ["general"],
            complexity_score=complexity,
            suggested_tier=tier,
            confidence=0.85,
        )

    def quick_enhance(self, prompt: str) -> str:
        return self.enhance(prompt).enhanced