import json
from typing import Any, cast

import openai
import structlog
from openai.types.chat import (
    ChatCompletion,
    ChatCompletionMessageParam,
    ChatCompletionNamedToolChoiceParam,
)
from pydantic import ValidationError

from spark_advisor_analyzer.ai.client import LlmClient
from spark_advisor_analyzer.ai.prompts import build_system_prompt, build_user_message
from spark_advisor_analyzer.ai.report_builder import build_advisor_report
from spark_advisor_analyzer.ai.tool_config import ANALYSIS_TOOL
from spark_advisor_models.config import AiSettings, Thresholds
from spark_advisor_models.model import (
    AdvisorReport,
    AnalysisToolInput,
    JobAnalysis,
    RuleResult,
    Severity,
)

logger = structlog.stdlib.get_logger(__name__)

_MAX_RETRIES = 3


class LlmAnalysisService:
    def __init__(self, client: LlmClient, ai: AiSettings, thresholds: Thresholds) -> None:
        self._client = client
        self._ai_settings = ai
        self._thresholds = thresholds
        self._prompt = build_system_prompt(self._thresholds)

    @property
    def model(self) -> str:
        return self._ai_settings.model

    def analyze(
        self,
        job: JobAnalysis,
        rule_results: list[RuleResult],
    ) -> AdvisorReport:
        user_msg: ChatCompletionMessageParam = {
            "role": "user",
            "content": build_user_message(job, rule_results, self._thresholds),
        }
        forced_tool: ChatCompletionNamedToolChoiceParam = {
            "type": "function",
            "function": {"name": "submit_analysis"},
        }

        for attempt in range(1, _MAX_RETRIES + 1):
            try:
                response = self._client.create_message(
                    model=self._ai_settings.model,
                    max_tokens=self._ai_settings.max_tokens,
                    system=self._prompt,
                    messages=[user_msg],
                    tools=[ANALYSIS_TOOL],
                    tool_choice=forced_tool,
                )
                raw = self._extract_tool_input(response)
                if not raw:
                    raise ValueError("Model returned empty tool arguments")
                parsed = AnalysisToolInput.model_validate(raw)
                return build_advisor_report(job.app_id, parsed, rule_results)
            except (ValueError, ValidationError) as exc:
                logger.warning("LLM tool call failed (attempt %d/%d): %s", attempt, _MAX_RETRIES, exc)
            except openai.RateLimitError:
                logger.warning("Rate limited (attempt %d/%d)", attempt, _MAX_RETRIES)

        logger.error("All %d LLM attempts failed, returning rules-only report", _MAX_RETRIES)
        return self._rules_only_report(job.app_id, rule_results)

    @staticmethod
    def _extract_tool_input(response: ChatCompletion) -> dict[str, Any]:
        tool_calls = response.choices[0].message.tool_calls
        if not tool_calls:
            raise ValueError("Model did not call submit_analysis tool")
        for tc in tool_calls:
            if hasattr(tc, "function") and tc.function.name == "submit_analysis":
                return cast("dict[str, Any]", json.loads(tc.function.arguments))
        raise ValueError("Model did not call submit_analysis tool")

    @staticmethod
    def _rules_only_report(app_id: str, rule_results: list[RuleResult]) -> AdvisorReport:
        worst = Severity.INFO
        for r in rule_results:
            if r.severity == Severity.CRITICAL:
                worst = Severity.CRITICAL
                break
            if r.severity == Severity.WARNING:
                worst = Severity.WARNING
        return AdvisorReport(
            app_id=app_id,
            summary="AI analysis unavailable — showing rule-based findings only.",
            severity=worst,
            rule_results=rule_results,
            recommendations=[],
        )
