import json
from typing import Any, cast

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

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,
)


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)

    def analyze(
        self,
        job: JobAnalysis,
        rule_results: list[RuleResult],
    ) -> AdvisorReport:
        user_message = build_user_message(job, rule_results, self._thresholds)
        user_msg: ChatCompletionMessageParam = {"role": "user", "content": user_message}
        forced_tool: ChatCompletionNamedToolChoiceParam = {
            "type": "function",
            "function": {"name": "submit_analysis"},
        }
        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_input = self._extract_tool_input(response)
        parsed = AnalysisToolInput.model_validate(raw_input)

        return build_advisor_report(job.app_id, parsed, 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")
