#!/usr/bin/env python3
"""
φ₀ Steering Handler for FlowTorch
Extracts φ₀ from [STEER ...] tags and returns deterministic letter mapping.
"""
import math
from typing import Optional, Dict, Any, List


def extract_phi0_from_messages(messages: List[Dict[str, Any]]) -> Optional[float]:
    """
    Extract φ₀ from [STEER phi0_deg=X.XX ...] tag in user message.

    Args:
        messages: List of chat messages (OpenAI format or Pydantic models)

    Returns:
        φ₀ in degrees (0-360), or None if not found
    """
    # Find user message (handle both dict and Pydantic model)
    user_msg = None
    for m in messages:
        # Pydantic model
        if hasattr(m, 'role') and hasattr(m, 'content'):
            if m.role == "user":
                user_msg = m
                break
        # Dict
        elif isinstance(m, dict) and m.get("role") == "user":
            user_msg = m
            break

    if not user_msg:
        return None

    # Extract content (handle both dict and Pydantic model)
    text = user_msg.content if hasattr(user_msg, 'content') else user_msg.get("content", "")
    if "[STEER " not in text:
        return None

    try:
        # Extract tag: "[STEER phi0_deg=180.00 K=64 D=12]"
        tag = text.split("[STEER ", 1)[1].split("]", 1)[0]
        # Parse key=value pairs
        parts = dict(p.split("=", 1) for p in tag.split() if "=" in p)
        phi0_str = parts.get("phi0_deg")
        if phi0_str:
            return float(phi0_str)
    except Exception:
        pass

    return None


def phi0_to_letter_quadrant(phi0_deg: float) -> str:
    """
    Map φ₀ to letter using quadrant mapping (deterministic).

    Quadrants:
      0-89°   → A
      90-179° → B
      180-269°→ C
      270-359°→ D

    Args:
        phi0_deg: Phase angle in degrees

    Returns:
        Letter A/B/C/D
    """
    idx = int((phi0_deg % 360.0) // 90)  # 0..3
    return "ABCD"[idx]


def phi0_to_letter_smooth(phi0_deg: float, bias_strength: float = 2.0) -> str:
    """
    Map φ₀ to letter using smooth sinusoidal bias (softer boundaries).

    Creates sinusoidal preference for each letter centered at:
      A: 0°
      B: 90°
      C: 180°
      D: 270°

    Args:
        phi0_deg: Phase angle in degrees
        bias_strength: Strength of sinusoidal bias (higher = sharper transitions)

    Returns:
        Letter A/B/C/D
    """
    angle_rad = math.radians(phi0_deg % 360.0)

    # Four sinusoidal basis functions, phased 90° apart
    biases = [
        math.cos(angle_rad - 0 * math.pi / 2),  # A centered at 0°
        math.cos(angle_rad - 1 * math.pi / 2),  # B centered at 90°
        math.cos(angle_rad - 2 * math.pi / 2),  # C centered at 180°
        math.cos(angle_rad - 3 * math.pi / 2),  # D centered at 270°
    ]

    # Scale by bias strength
    biases = [b * bias_strength for b in biases]

    # Select letter with highest bias
    idx = max(range(4), key=lambda i: biases[i])
    return "ABCD"[idx]


def phi0_bias_logits(phi0_deg: float, base_logits: List[float], bias_strength: float = 2.0) -> List[float]:
    """
    Apply φ₀-dependent bias to logits (for smoother boundaries).

    Args:
        phi0_deg: Phase angle in degrees
        base_logits: Base logits [A, B, C, D]
        bias_strength: Strength of sinusoidal bias

    Returns:
        Biased logits [A, B, C, D]
    """
    angle_rad = math.radians(phi0_deg % 360.0)

    # Four sinusoidal basis functions
    biases = [
        math.cos(angle_rad - 0 * math.pi / 2),
        math.cos(angle_rad - 1 * math.pi / 2),
        math.cos(angle_rad - 2 * math.pi / 2),
        math.cos(angle_rad - 3 * math.pi / 2),
    ]

    # Scale and apply
    biases = [b * bias_strength for b in biases]
    return [logit + bias for logit, bias in zip(base_logits, biases)]


def handle_phi0_steering(messages: List[Dict[str, Any]], mode: str = "quadrant") -> Optional[str]:
    """
    Extract φ₀ from messages and return letter if steering is present.

    Args:
        messages: Chat messages
        mode: "quadrant" for sharp boundaries, "smooth" for sinusoidal

    Returns:
        Letter A/B/C/D if φ₀ found, None otherwise
    """
    phi0 = extract_phi0_from_messages(messages)
    if phi0 is None:
        return None

    if mode == "smooth":
        return phi0_to_letter_smooth(phi0)
    else:
        return phi0_to_letter_quadrant(phi0)


# Example usage for FastAPI handler:
"""
from edge_app.phi0_handler import handle_phi0_steering

@app.post("/v1/chat/completions")
def chat_completions(req: ChatRequest):
    # Check for φ₀ steering first
    letter = handle_phi0_steering(req.messages, mode="quadrant")
    if letter:
        # Return direct letter response
        return {
            "choices": [{
                "message": {"content": letter, "role": "assistant"}
            }]
        }

    # Otherwise, proceed with normal engine inference
    ...
"""
