from dataclasses import dataclass
from typing import Any, get_origin, get_args

from pydantic import BaseModel


@dataclass
class Usage:
    """Token usage statistics."""

    input_tokens: int = 0
    cached_input_tokens: int = 0
    output_tokens: int = 0

    @property
    def total_tokens(self) -> int:
        """Total tokens used (input + output)."""
        return self.input_tokens + self.output_tokens

    def __repr__(self) -> str:
        return (
            f"Usage(input={self.input_tokens}, "
            f"cached={self.cached_input_tokens}, "
            f"output={self.output_tokens}, "
            f"total={self.total_tokens})"
        )


def is_pydantic_model(type_hint: Any) -> bool:
    """Check if a type hint is a Pydantic model."""
    try:
        return isinstance(type_hint, type) and issubclass(type_hint, BaseModel)
    except TypeError:
        return False


def get_json_schema(pydantic_model: type[BaseModel]) -> dict[str, Any]:
    """Generate JSON schema from Pydantic model.

    Sets additionalProperties to false to match Codex behavior.
    """
    schema = pydantic_model.model_json_schema()

    # Recursively set additionalProperties: false for all objects
    def set_additional_properties(obj: dict[str, Any]) -> None:
        if isinstance(obj, dict):
            if obj.get("type") == "object":
                obj["additionalProperties"] = False
            for value in obj.values():
                if isinstance(value, dict):
                    set_additional_properties(value)
                elif isinstance(value, list):
                    for item in value:
                        if isinstance(item, dict):
                            set_additional_properties(item)

    set_additional_properties(schema)
    return schema


def parse_json_response(response: str, pydantic_model: type[BaseModel]) -> BaseModel:
    """Parse JSON response into Pydantic model.

    Raises:
        json.JSONDecodeError: If response is not valid JSON
        pydantic.ValidationError: If JSON doesn't match model schema
    """
    return pydantic_model.model_validate_json(response)


def extract_text_from_response(response: str) -> str:
    """Extract plain text from response.

    For string outputs, we just return the text as-is.
    Claude Code may wrap responses in markdown or other formatting.
    """
    return response.strip()