from __future__ import annotations

import json
import re
from typing import Any, Protocol

from tenacity import retry, stop_after_attempt, wait_exponential


class LLMService(Protocol):
    def complete(self, system: str, user: str, max_tokens: int = 4096) -> str: ...

    def complete_json(self, system: str, user: str, max_tokens: int = 4096) -> Any: ...


def parse_json_loose(text: str) -> Any:
    text = text.strip()
    fence = re.search(r"```(?:json)?\s*([\s\S]*?)```", text)
    if fence:
        text = fence.group(1).strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        start = text.find("{")
        end = text.rfind("}")
        if start >= 0 and end > start:
            return json.loads(text[start : end + 1])
        start = text.find("[")
        end = text.rfind("]")
        if start >= 0 and end > start:
            return json.loads(text[start : end + 1])
        raise


class ClaudeLLM:
    """Anthropic Claude — used for ideation + platform copy."""

    def __init__(self, api_key: str, model: str):
        import anthropic

        self.client = anthropic.Anthropic(api_key=api_key)
        self.model = model

    @retry(stop=stop_after_attempt(3), wait=wait_exponential(min=2, max=20), reraise=True)
    def complete(self, system: str, user: str, max_tokens: int = 4096) -> str:
        msg = self.client.messages.create(
            model=self.model,
            max_tokens=max_tokens,
            system=system,
            messages=[{"role": "user", "content": user}],
        )
        parts = [b.text for b in msg.content if getattr(b, "type", None) == "text"]
        return "\n".join(parts).strip()

    def complete_json(self, system: str, user: str, max_tokens: int = 4096) -> Any:
        raw = self.complete(
            system=system + "\n\nRespond with valid JSON only. No markdown fences.",
            user=user,
            max_tokens=max_tokens,
        )
        return parse_json_loose(raw)


class OpenAILLM:
    """OpenAI Chat Completions — optional fallback for ideation."""

    def __init__(self, api_key: str, model: str):
        from openai import OpenAI

        self.client = OpenAI(api_key=api_key)
        self.model = model

    @retry(stop=stop_after_attempt(3), wait=wait_exponential(min=2, max=20), reraise=True)
    def complete(self, system: str, user: str, max_tokens: int = 4096) -> str:
        response = self.client.chat.completions.create(
            model=self.model,
            max_tokens=max_tokens,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
        )
        return (response.choices[0].message.content or "").strip()

    def complete_json(self, system: str, user: str, max_tokens: int = 4096) -> Any:
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                max_tokens=max_tokens,
                response_format={"type": "json_object"},
                messages=[
                    {"role": "system", "content": system + "\n\nRespond with valid JSON only."},
                    {"role": "user", "content": user},
                ],
            )
            raw = (response.choices[0].message.content or "").strip()
            return parse_json_loose(raw)
        except Exception:
            raw = self.complete(
                system=system + "\n\nRespond with valid JSON only. No markdown fences.",
                user=user,
                max_tokens=max_tokens,
            )
            return parse_json_loose(raw)


def create_llm(*, provider: str, anthropic_key: str, claude_model: str, openai_key: str, openai_model: str):
    provider = (provider or "openai").strip().lower()
    if provider == "openai":
        if not openai_key:
            raise RuntimeError("LLM_PROVIDER=openai but OPENAI_API_KEY is missing")
        return OpenAILLM(openai_key, openai_model)
    if provider == "claude":
        if not anthropic_key:
            raise RuntimeError("LLM_PROVIDER=claude but ANTHROPIC_API_KEY is missing")
        return ClaudeLLM(anthropic_key, claude_model)
    raise RuntimeError(f"Unknown LLM_PROVIDER: {provider} (use claude or openai)")