"""
TokioAI Engine -- Output Verifier & Retry Logic
Validates model responses and orchestrates retries on failure.

The verifier catches:
1. Invalid JSON tool calls
2. Wrong tool names
3. Missing required parameters
4. Malformed output (e.g., natural language when tool expected)
5. Verbosity (strips filler from responses)
"""

import json
import re
import time
from typing import Any, Dict, List, Optional, Tuple

from .schemas import validate_tool_call, extract_tool_calls_from_text, VALID_TOOL_NAMES


class VerificationResult:
    """Result of verifying model output."""

    def __init__(
        self,
        valid: bool,
        output_type: str,  # "tool_call", "text", "error"
        tool_calls: List[Dict] = None,
        text: str = None,
        error: str = None,
        raw_output: str = None,
    ):
        self.valid = valid
        self.output_type = output_type
        self.tool_calls = tool_calls or []
        self.text = text or ""
        self.error = error or ""
        self.raw_output = raw_output or ""

    def __repr__(self):
        if self.valid:
            if self.output_type == "tool_call":
                names = [tc["name"] for tc in self.tool_calls]
                return f"VerificationResult(valid=True, tools={names})"
            return f"VerificationResult(valid=True, type={self.output_type})"
        return f"VerificationResult(valid=False, error={self.error!r})"


class Verifier:
    """
    Validates model output and provides structured results.
    """

    def __init__(self, strict: bool = False):
        """
        Args:
            strict: If True, rejects ambiguous outputs. If False, tries harder to parse.
        """
        self.strict = strict

    def verify(self, model_output: str, expected_intent: str = None) -> VerificationResult:
        """
        Verify model output.
        
        Args:
            model_output: Raw text from the model
            expected_intent: What the router classified (tool_call, text, etc.)
            
        Returns:
            VerificationResult with parsed tool calls or cleaned text
        """
        if not model_output or not model_output.strip():
            return VerificationResult(
                valid=False,
                output_type="error",
                error="Empty model output",
                raw_output=model_output,
            )

        raw = model_output.strip()

        # Step 1: Try to extract tool calls
        tool_calls = extract_tool_calls_from_text(raw)

        if tool_calls:
            # Validate each tool call
            valid_calls = []
            errors = []

            for tc in tool_calls:
                is_valid, err = validate_tool_call(tc["name"], tc["arguments"])
                if is_valid:
                    valid_calls.append(tc)
                else:
                    errors.append(err)

            if valid_calls:
                return VerificationResult(
                    valid=True,
                    output_type="tool_call",
                    tool_calls=valid_calls,
                    raw_output=raw,
                )
            elif errors:
                return VerificationResult(
                    valid=False,
                    output_type="error",
                    error="; ".join(errors),
                    raw_output=raw,
                )

        # Step 2: If expected tool_call but got text, that might be an error
        if expected_intent == "tool_call":
            # Check if the model is TRYING to make a tool call but failing
            if any(name in raw.lower() for name in VALID_TOOL_NAMES):
                return VerificationResult(
                    valid=False,
                    output_type="error",
                    error="Model mentioned a tool name but didn't format it as a valid JSON tool call",
                    raw_output=raw,
                )

        # Step 3: It's a text response -- clean it up
        cleaned = self._clean_text(raw)
        return VerificationResult(
            valid=True,
            output_type="text",
            text=cleaned,
            raw_output=raw,
        )

    def _clean_text(self, text: str) -> str:
        """Clean up verbose model output."""
        # Remove thinking tags (Qwen3 /think blocks)
        text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
        # Remove empty lines at start/end
        text = text.strip()
        # Remove excessive newlines
        text = re.sub(r'\n{3,}', '\n\n', text)
        return text


class RetryStrategy:
    """
    Manages retry logic when the verifier rejects model output.
    Each retry gets a more explicit prompt to guide the model.
    """

    def __init__(self, max_retries: int = 2):
        self.max_retries = max_retries

    def should_retry(self, attempt: int, result: VerificationResult) -> bool:
        """Should we retry after a failed verification?"""
        if result.valid:
            return False
        if attempt >= self.max_retries:
            return False
        # Don't retry empty outputs (model is broken)
        if not result.raw_output.strip():
            return False
        return True

    def get_retry_message(self, attempt: int, result: VerificationResult) -> str:
        """Generate a retry prompt based on what went wrong."""
        if attempt == 0:
            # First retry: gentle correction
            return (
                f"Your previous response was not valid: {result.error}\n"
                f"Please respond with ONLY a JSON object like this:\n"
                f'{{"name": "tool_name", "arguments": {{"param": "value"}}}}\n'
                f"No explanation, no markdown, just the JSON."
            )
        else:
            # Second retry: very explicit
            return (
                f"ERROR: {result.error}\n"
                f"You MUST respond with valid JSON only. "
                f"Valid tool names: {', '.join(sorted(VALID_TOOL_NAMES))}.\n"
                f"Format: {{\"name\": \"execute_local\", \"arguments\": {{\"command\": \"your command\"}}}}\n"
                f"Respond with NOTHING except the JSON object."
            )


def verify_and_extract(model_output: str, expected_intent: str = None) -> VerificationResult:
    """Convenience function: verify model output with default settings."""
    return Verifier().verify(model_output, expected_intent)