from __future__ import annotations

import json
from functools import lru_cache
from typing import Any, Dict, List, Optional

from google.api_core.exceptions import GoogleAPICallError, InvalidArgument
from jsonschema import Draft7Validator, ValidationError
from vertexai import init as vertex_init
from vertexai.generative_models import (
    FinishReason,
    Content,
    GenerationConfig,
    GenerativeModel,
    HarmBlockThreshold,
    HarmCategory,
    Part,
    SafetySetting,
)

from app.core.config import get_settings


class LLMContentFilteredError(Exception):
    def __init__(self, category: Optional[str] = None):
        self.category = category or "content_filtered"
        super().__init__(self.category)


def _format_messages(messages: List[Dict[str, Any]]) -> List[Content]:
    contents: List[Content] = []
    for message in messages:
        role = message.get("role", "user")
        if role == "assistant":
            vertex_role = "model"
        elif role == "system":
            vertex_role = "user"
        else:
            vertex_role = role
        contents.append(
            Content(
                role=vertex_role,
                parts=[Part.from_text(str(message.get("content", "")))],
            )
        )
    return contents


@lru_cache
def _get_model() -> GenerativeModel:
    settings = get_settings()
    vertex_init(project=settings.vertex_ai.project_id, location=settings.vertex_ai.location)
    return GenerativeModel(settings.vertex_ai.model)


def invoke_chat(
    messages: List[Dict[str, Any]],
    response_format: Optional[Dict[str, Any]] = None,
    response_schema: Optional[Dict[str, Any]] = None,
    **kwargs: Any,
) -> Dict[str, Any]:
    model = _get_model()
    contents = _format_messages(messages)

    temperature = kwargs.get("temperature", 0.2)
    max_tokens = kwargs.get("max_tokens")

    generation_config = GenerationConfig(
        temperature=temperature,
        max_output_tokens=max_tokens,
        response_mime_type="application/json",
    )

    schema_validator: Optional[Draft7Validator] = None
    if response_schema is not None:
        schema_validator = Draft7Validator(response_schema)

    safety_settings = [
        SafetySetting(
            category=HarmCategory.HARM_CATEGORY_HATE_SPEECH,
            threshold=HarmBlockThreshold.BLOCK_NONE,
        ),
        SafetySetting(
            category=HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
            threshold=HarmBlockThreshold.BLOCK_NONE,
        ),
    SafetySetting(
        category=HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
        threshold=HarmBlockThreshold.BLOCK_NONE,
    ),
        SafetySetting(
            category=HarmCategory.HARM_CATEGORY_HARASSMENT,
            threshold=HarmBlockThreshold.BLOCK_NONE,
        ),
    ]

    try:
        response = model.generate_content(
            contents,
            generation_config=generation_config,
            safety_settings=safety_settings,
        )
    except (GoogleAPICallError, InvalidArgument) as exc:  # pragma: no cover
        message = str(exc)
        if "HARM_CATEGORY" in message:
            category = message.split("HARM_CATEGORY")[-1].strip(" :")
            raise LLMContentFilteredError(category) from exc
        raise RuntimeError("Vertex AI call failed") from exc

    candidates = getattr(response, "candidates", []) or []
    if not candidates:
        raise RuntimeError("Vertex AI returned no candidates")

    candidate = candidates[0]
    finish_reason = getattr(candidate, "finish_reason", None)
    if finish_reason == FinishReason.SAFETY:
        safety_ratings = getattr(candidate, "safety_ratings", []) or []
        blocked_category: Optional[str] = None
        for rating in safety_ratings:
            if getattr(rating, "blocked", False):
                category = getattr(rating, "category", None)
                blocked_category = getattr(category, "name", str(category))
                break
        raise LLMContentFilteredError(blocked_category)

    text_response = getattr(candidate, "text", None)
    if not text_response:
        text_response = getattr(candidate, "output_text", None)
    if not text_response and getattr(candidate, "content", None):
        parts = getattr(candidate.content, "parts", [])
        text_response = "".join(part.text for part in parts if getattr(part, "text", None))
    if not text_response:
        text_response = getattr(response, "text", None)
    if not text_response:
        raise RuntimeError("Vertex AI returned empty response")

    try:
        data = json.loads(text_response)
        if schema_validator is not None:
            errors = sorted(schema_validator.iter_errors(data), key=lambda err: list(err.path))
            if errors:
                message = "; ".join(err.message for err in errors)
                raise ValueError(f"LLM response failed schema validation: {message}")
        return data
    except json.JSONDecodeError:
        if response_format and response_format.get("type") == "json_object":
            raise
        return {"output": text_response}
