# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Structured Outputs for Familiar

Provides type-safe parsing and validation of LLM responses using Pydantic.

This eliminates fragile regex-based JSON extraction and provides:
- Strong typing for LLM outputs
- Automatic validation with clear error messages
- Retry-friendly error handling
- Reusable schema definitions

Usage:
    from familiar.core.structured import parse_json, PlanOutput, ReflectionOutput

    # Parse with validation
    result = parse_json(response_text, PlanOutput)
    if result.success:
        plan = result.value
    else:
        print(f"Parse error: {result.error}")

    # Or raise on failure
    plan = parse_json(response_text, PlanOutput).unwrap()
"""

import json
import logging
import re
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Generic, List, Optional, Type, TypeVar

# Try to import Pydantic v2, fall back to v1 patterns if needed
try:
    from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator

    PYDANTIC_V2 = hasattr(BaseModel, "model_validate")
except ImportError:
    # Provide a minimal fallback if Pydantic isn't installed
    PYDANTIC_V2 = False
    BaseModel = None
    ValidationError = Exception

    def Field(*a, **kw):
        return None

    def field_validator(*a, **kw):
        return lambda f: f

    ConfigDict = None

logger = logging.getLogger(__name__)

T = TypeVar("T")


# ============================================================
# RESULT TYPE (for parse operations)
# ============================================================


@dataclass
class ParseResult(Generic[T]):
    """
    Result of a parse operation.

    Either contains a validated value or an error message.
    """

    success: bool
    value: Optional[T] = None
    error: Optional[str] = None
    raw_json: Optional[dict] = None  # The extracted JSON before validation

    def unwrap(self) -> T:
        """Return value or raise exception."""
        if self.success and self.value is not None:
            return self.value
        raise ValueError(self.error or "Parse failed with no error message")

    def unwrap_or(self, default: T) -> T:
        """Return value or default."""
        if self.success and self.value is not None:
            return self.value
        return default

    def map(self, func):
        """Apply function to value if successful."""
        if self.success and self.value is not None:
            return ParseResult(success=True, value=func(self.value))
        return self


# ============================================================
# JSON EXTRACTION
# ============================================================


def extract_json(text: str) -> Optional[dict]:
    """
    Extract JSON from LLM response text.

    Handles common patterns:
    - Raw JSON
    - JSON in markdown code blocks (```json ... ```)
    - JSON mixed with prose
    - Multiple JSON objects (returns first valid one)

    Returns None if no valid JSON found.
    """
    if not text or not text.strip():
        return None

    text = text.strip()

    # Strategy 1: Try parsing the whole text as JSON
    try:
        return json.loads(text)
    except json.JSONDecodeError as e:
        logger.debug("JSON parse error suppressed: %s", e)
    # Strategy 2: Extract from markdown code blocks
    code_block_patterns = [
        r"```json\s*([\s\S]*?)\s*```",  # ```json ... ```
        r"```\s*([\s\S]*?)\s*```",  # ``` ... ```
    ]

    for pattern in code_block_patterns:
        matches = re.findall(pattern, text)
        for match in matches:
            try:
                return json.loads(match.strip())
            except json.JSONDecodeError as e:
                logger.debug("JSON parse error suppressed: %s", e)
                continue
    # Strategy 3: Find JSON object boundaries
    # Look for outermost { } pair
    brace_start = text.find("{")
    if brace_start != -1:
        # Find matching closing brace
        depth = 0
        in_string = False
        escape_next = False

        for i, char in enumerate(text[brace_start:], start=brace_start):
            if escape_next:
                escape_next = False
                continue

            if char == "\\" and in_string:
                escape_next = True
                continue

            if char == '"' and not escape_next:
                in_string = not in_string
                continue

            if not in_string:
                if char == "{":
                    depth += 1
                elif char == "}":
                    depth -= 1
                    if depth == 0:
                        try:
                            return json.loads(text[brace_start : i + 1])
                        except json.JSONDecodeError:
                            break

    # Strategy 4: Find JSON array boundaries
    bracket_start = text.find("[")
    if bracket_start != -1:
        depth = 0
        in_string = False
        escape_next = False

        for i, char in enumerate(text[bracket_start:], start=bracket_start):
            if escape_next:
                escape_next = False
                continue

            if char == "\\" and in_string:
                escape_next = True
                continue

            if char == '"' and not escape_next:
                in_string = not in_string
                continue

            if not in_string:
                if char == "[":
                    depth += 1
                elif char == "]":
                    depth -= 1
                    if depth == 0:
                        try:
                            return json.loads(text[bracket_start : i + 1])
                        except json.JSONDecodeError:
                            break

    return None


def parse_json(text: str, model: Optional[Type[T]] = None, strict: bool = False) -> ParseResult[T]:
    """
    Parse JSON from text and optionally validate against a Pydantic model.

    Args:
        text: The raw text (potentially containing JSON)
        model: Optional Pydantic model class to validate against
        strict: If True, fail on extra fields; if False, ignore them

    Returns:
        ParseResult with either the validated model or an error

    Examples:
        # Just extract JSON
        result = parse_json('{"key": "value"}')
        data = result.unwrap()  # {'key': 'value'}

        # Validate against model
        result = parse_json(response, PlanOutput)
        plan = result.unwrap()  # PlanOutput instance
    """
    # Extract JSON
    raw_json = extract_json(text)

    if raw_json is None:
        return ParseResult(success=False, error=f"No valid JSON found in text: {text[:200]}...")

    # If no model specified, return raw dict
    if model is None:
        return ParseResult(success=True, value=raw_json, raw_json=raw_json)

    # Validate against Pydantic model
    if BaseModel is None:
        return ParseResult(success=False, error="Pydantic not available.")

    try:
        if PYDANTIC_V2:
            validated = model.model_validate(raw_json, strict=strict)
        else:
            validated = model.parse_obj(raw_json)

        return ParseResult(success=True, value=validated, raw_json=raw_json)

    except ValidationError as e:
        error_msg = _format_validation_error(e)
        return ParseResult(
            success=False, error=f"Validation failed: {error_msg}", raw_json=raw_json
        )
    except Exception as e:
        return ParseResult(
            success=False, error=f"Unexpected error during validation: {e}", raw_json=raw_json
        )


def _format_validation_error(error: ValidationError) -> str:
    """Format Pydantic validation error for readability."""
    if PYDANTIC_V2:
        errors = error.errors()
    else:
        errors = error.errors()

    messages = []
    for err in errors:
        loc = ".".join(str(x) for x in err.get("loc", []))
        msg = err.get("msg", "Unknown error")
        messages.append(f"{loc}: {msg}")

    return "; ".join(messages)


# ============================================================
# SCHEMA DEFINITIONS (Pydantic Models)
# ============================================================

if BaseModel is not None:

    class StepActionType(str, Enum):
        """Valid action types for plan steps."""

        TOOL = "tool"
        CHAT = "chat"
        CONDITION = "condition"
        PARALLEL = "parallel"
        HUMAN = "human"

    class StepConfig(BaseModel):
        """Configuration for a step action."""

        model_config = ConfigDict(extra="allow")  # Allow additional fields

        # For tool actions
        tool: Optional[str] = None
        input: Optional[Dict[str, Any]] = Field(default_factory=dict)

        # For chat actions
        prompt: Optional[str] = None

        # For parallel actions
        steps: Optional[List[str]] = None

    class PlanStep(BaseModel):
        """A single step in a generated plan."""

        id: str
        description: str
        action_type: str = Field(description="One of: tool, chat, condition, parallel, human")
        action_config: Dict[str, Any] = Field(default_factory=dict)
        depends_on: List[str] = Field(default_factory=list)
        is_critical: bool = True
        estimated_duration: Optional[int] = None

        # For conditional steps
        condition: Optional[str] = None
        if_true: Optional[str] = None
        if_false: Optional[str] = None

        @field_validator("action_type")
        @classmethod
        def validate_action_type(cls, v):
            valid_types = {"tool", "chat", "condition", "parallel", "human"}
            if v.lower() not in valid_types:
                raise ValueError(f"action_type must be one of {valid_types}")
            return v.lower()

    class PlanOutput(BaseModel):
        """Output schema for plan generation."""

        description: str = ""
        steps: List[PlanStep]

        @field_validator("steps")
        @classmethod
        def validate_steps_not_empty(cls, v):
            if not v:
                raise ValueError("Plan must have at least one step")
            return v

    class ReflectionAction(str, Enum):
        """Valid actions for reflection/self-correction."""

        RETRY = "RETRY"
        MODIFY = "MODIFY"
        SKIP = "SKIP"
        REPLAN = "REPLAN"
        ABORT = "ABORT"

    class ReflectionOutput(BaseModel):
        """Output schema for failure reflection."""

        action: str = Field(description="One of: RETRY, MODIFY, SKIP, REPLAN, ABORT")
        reason: str
        modifications: Dict[str, Any] = Field(default_factory=dict)
        new_steps: List[Dict[str, Any]] = Field(default_factory=list)

        @field_validator("action")
        @classmethod
        def validate_action(cls, v):
            valid_actions = {"RETRY", "MODIFY", "SKIP", "REPLAN", "ABORT"}
            upper_v = v.upper()
            if upper_v not in valid_actions:
                raise ValueError(f"action must be one of {valid_actions}")
            return upper_v

    class ToolCallOutput(BaseModel):
        """Generic schema for tool call responses that return JSON."""

        success: bool = True
        data: Optional[Dict[str, Any]] = None
        error: Optional[str] = None
        message: Optional[str] = None

    class ExtractionOutput(BaseModel):
        """Schema for entity/information extraction tasks."""

        entities: List[Dict[str, Any]] = Field(default_factory=list)
        relationships: List[Dict[str, Any]] = Field(default_factory=list)
        summary: Optional[str] = None
        confidence: float = Field(default=1.0, ge=0.0, le=1.0)

    class ClassificationOutput(BaseModel):
        """Schema for classification tasks."""

        label: str
        confidence: float = Field(ge=0.0, le=1.0)
        reasoning: Optional[str] = None
        alternatives: List[Dict[str, Any]] = Field(default_factory=list)

    class SentimentOutput(BaseModel):
        """Schema for sentiment analysis."""

        sentiment: str = Field(description="One of: positive, negative, neutral, mixed")
        score: float = Field(ge=-1.0, le=1.0)
        aspects: List[Dict[str, Any]] = Field(default_factory=list)

        @field_validator("sentiment")
        @classmethod
        def validate_sentiment(cls, v):
            valid = {"positive", "negative", "neutral", "mixed"}
            if v.lower() not in valid:
                raise ValueError(f"sentiment must be one of {valid}")
            return v.lower()

else:
    # Fallback when Pydantic not available
    PlanOutput = None
    PlanStep = None
    ReflectionOutput = None
    ToolCallOutput = None
    ExtractionOutput = None
    ClassificationOutput = None
    SentimentOutput = None


# ============================================================
# PROMPT HELPERS
# ============================================================


def schema_to_prompt(model: Type[BaseModel]) -> str:
    """
    Convert a Pydantic model to a prompt-friendly schema description.

    This helps guide the LLM to produce valid output.
    """
    if BaseModel is None or model is None:
        return ""

    if PYDANTIC_V2:
        schema = model.model_json_schema()
    else:
        schema = model.schema()

    return json.dumps(schema, indent=2)


def create_json_prompt(
    model: Type[BaseModel], task_description: str, examples: Optional[List[dict]] = None
) -> str:
    """
    Create a prompt that instructs the LLM to output valid JSON.

    Args:
        model: Pydantic model defining the expected output
        task_description: Description of what to do
        examples: Optional list of example outputs

    Returns:
        A formatted prompt string
    """
    schema_str = schema_to_prompt(model)

    prompt_parts = [
        task_description,
        "",
        "Output your response as valid JSON matching this schema:",
        "```json",
        schema_str,
        "```",
    ]

    if examples:
        prompt_parts.extend(
            [
                "",
                "Examples of valid outputs:",
            ]
        )
        for i, example in enumerate(examples, 1):
            prompt_parts.extend(
                [
                    f"Example {i}:",
                    "```json",
                    json.dumps(example, indent=2),
                    "```",
                ]
            )

    prompt_parts.extend(
        [
            "",
            "Important: Output ONLY valid JSON, no additional text or explanation.",
        ]
    )

    return "\n".join(prompt_parts)


# ============================================================
# CONVENIENCE FUNCTIONS
# ============================================================


def parse_plan(text: str) -> ParseResult[PlanOutput]:
    """Parse a plan generation response."""
    if PlanOutput is None:
        return ParseResult(success=False, error="Pydantic not installed")
    return parse_json(text, PlanOutput)


def parse_reflection(text: str) -> ParseResult[ReflectionOutput]:
    """Parse a reflection/self-correction response."""
    if ReflectionOutput is None:
        return ParseResult(success=False, error="Pydantic not installed")
    return parse_json(text, ReflectionOutput)


def parse_tool_response(text: str) -> ParseResult[ToolCallOutput]:
    """Parse a tool call response."""
    if ToolCallOutput is None:
        return ParseResult(success=False, error="Pydantic not installed")
    return parse_json(text, ToolCallOutput)
