from __future__ import annotations

import ast
import json
import re
from abc import ABC, abstractmethod
from enum import Enum, auto
from typing import Dict, List, Literal, Type, overload

from openai import APIError, Client, omit
from openai.types.chat import ChatCompletionMessageParam
from pydantic import BaseModel, ValidationError

from srs2std.log import LoggingMixin


class Prompt(BaseModel):
    system: str
    user: str


class Response[T](BaseModel):
    content: T | None
    finish_reason: FinishReason


class FinishReason(Enum):
    STOP = auto()
    FILTERED = auto()
    MAX_RETRIES = auto()
    PARSE_FAILED = auto()


class LLM(ABC, LoggingMixin):
    def __init__(self):
        super().__init__()

    @property
    @abstractmethod
    def client(self) -> Client:
        pass

    @property
    @abstractmethod
    def model(self) -> str:
        pass

    @property
    @abstractmethod
    def max_completion_tokens(self) -> int:
        pass

    @property
    @abstractmethod
    def max_retries(self) -> int:
        pass

    @property
    @abstractmethod
    def temperature(self) -> float:
        pass

    @property
    def reasoning_effort(self) -> Literal["low", "medium", "high"] | None:
        return "medium"

    @overload
    def ask(self, prompt: Prompt) -> Response[str]: ...

    @overload
    def ask[T: BaseModel](self, prompt: Prompt, schema: Type[T]) -> Response[T]: ...

    def ask[T: BaseModel](
        self,
        prompt: Prompt,
        schema: Type[T] | None = None,
    ) -> Response[T] | Response[str]:
        if schema is None or schema is str:
            return self._create_completion(prompt)
        else:
            return self._parse_completion(prompt, schema)

    def _create_completion(self, prompt: Prompt) -> Response[str]:
        messages: List[ChatCompletionMessageParam] = [
            {"role": "system", "content": prompt.system},
            {"role": "user", "content": prompt.user},
        ]

        completions = []
        for _ in range(self.max_retries):
            self.log_info(messages)

            completion = None
            try:
                completion = self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    max_completion_tokens=self.max_completion_tokens,
                    temperature=self.temperature,
                    reasoning_effort=self.reasoning_effort or omit,
                )
            except APIError as e:
                self.log_error(e)
                continue

            message = completion.choices[0].message
            finish_reason = completion.choices[0].finish_reason

            if finish_reason == "content_filter":
                self.log_error(message.model_dump_json())
                return Response(content=None, finish_reason=FinishReason.FILTERED)
            self.log_info(message.model_dump_json())

            completions.append(message.content)

            if finish_reason == "stop":
                content = " ".join(completions)
                return Response(content=content, finish_reason=FinishReason.STOP)

            messages.extend(
                [
                    {"role": "assistant", "content": message.content or ""},
                    {"role": "user", "content": "Continue generating."},
                ]
            )

        else:
            return Response(content=None, finish_reason=FinishReason.MAX_RETRIES)

    def _parse_completion[T: BaseModel](
        self, prompt: Prompt, schema: Type[T]
    ) -> Response[T]:
        messages: List[ChatCompletionMessageParam] = [
            {"role": "system", "content": prompt.system},
            {"role": "user", "content": prompt.user},
        ]

        completions = []
        for _ in range(self.max_retries):
            self.log_info(messages)

            completion = None
            try:
                completion = self.client.chat.completions.parse(
                    model=self.model,
                    messages=messages,
                    max_completion_tokens=self.max_completion_tokens,
                    temperature=self.temperature,
                    reasoning_effort=self.reasoning_effort or omit,
                    response_format=schema,  # specify structured output
                )
            except APIError as e:
                self.log_error(e)
                continue

            message = completion.choices[0].message
            finish_reason = completion.choices[0].finish_reason

            if finish_reason == "content_filter":
                self.log_error(message.model_dump_json())
                return Response(content=None, finish_reason=FinishReason.FILTERED)
            self.log_info(message.model_dump_json())

            parsed = message.parsed
            if parsed is not None:
                return Response(content=parsed, finish_reason=FinishReason.STOP)

            completions.append(message.content)

            if finish_reason == "stop":
                content = " ".join(completions)
                return fallback_parse(content, schema)

            messages.extend(
                [
                    {"role": "assistant", "content": message.content or ""},
                    {"role": "user", "content": "Continue generating."},
                ]
            )

        else:
            return Response(content=None, finish_reason=FinishReason.MAX_RETRIES)


def fallback_parse[T: BaseModel](text_content: str, schema: Type[T]) -> Response[T]:
    match = re.search(
        pattern=r"(?:```(?:json)?)?\s*(\{.+\})\s*(?:```)?",
        string=text_content,
        flags=re.DOTALL,
    )
    cleaned_text = match.group(1) if match else text_content

    try:
        parsed = json.loads(cleaned_text)
    except json.JSONDecodeError:
        try:
            parsed = ast.literal_eval(cleaned_text)
        except Exception:
            return Response(content=None, finish_reason=FinishReason.PARSE_FAILED)

    if not isinstance(parsed, Dict):
        return Response(content=None, finish_reason=FinishReason.PARSE_FAILED)

    try:
        content = schema(**parsed)
        return Response(content=content, finish_reason=FinishReason.STOP)
    except ValidationError:
        return Response(content=None, finish_reason=FinishReason.PARSE_FAILED)
