"""Main client for generating WhatsApp Flows and Message Templates using LLMs."""

from __future__ import annotations

import json
import os
import re
from typing import Any, Callable, Optional

from openai import OpenAI
from pydantic import ValidationError

from roovy_sdk.docs import get_enhanced_prompt
from roovy_sdk.prompt import SYSTEM_PROMPT
from roovy_sdk.schema import WhatsAppFlow, validate_flow
from roovy_sdk.template_prompt import TEMPLATE_DOCS, TEMPLATE_SYSTEM_PROMPT
from roovy_sdk.template_schema import TemplateResult, validate_template_result


def _clean_json_response(content: str) -> str:
    """Clean LLM output to extract the JSON object.

    - Strip markdown ```json fences
    - Take the substring from first '{' to last '}'
    """
    json_str = content.strip()

    # Remove markdown code fences
    if json_str.startswith("```"):
        json_str = re.sub(r"```json?\s*", "", json_str, flags=re.IGNORECASE)
        json_str = re.sub(r"```\s*$", "", json_str).strip()

    # Trim to outermost JSON object
    start_idx = json_str.find("{")
    end_idx = json_str.rfind("}")
    if start_idx != -1 and end_idx != -1 and end_idx >= start_idx:
        json_str = json_str[start_idx : end_idx + 1]

    return json_str


def _format_validation_errors(error: ValidationError) -> str:
    """Format Pydantic validation errors into a readable string."""
    lines = []
    for e in error.errors():
        loc = ".".join(str(p) for p in e.get("loc", []))
        msg = e.get("msg", "validation error")
        lines.append(f"{loc}: {msg}")
    return "\n".join(lines)


class RoovyClient:
    """Client for generating WhatsApp Flows using LLMs.

    Args:
        api_key: API key for the LLM provider. Falls back to ROOVY_API_KEY env var.
        base_url: Base URL for the OpenAI-compatible API. Falls back to ROOVY_BASE_URL env var.
        model: Model name to use. Falls back to ROOVY_MODEL env var.

    Example:
        >>> client = RoovyClient(
        ...     api_key="sk-...",
        ...     base_url="https://api.openai.com/v1",
        ...     model="gpt-4o",
        ... )
        >>> flow = client.generate("Create a feedback form")
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        model: Optional[str] = None,
        *,
        use_docs: bool = True,
        custom_docs_path: Optional[str] = None,
        custom_docs: Optional[str] = None,
    ) -> None:
        self.api_key = api_key or os.getenv("ROOVY_API_KEY", "")
        self.base_url = base_url or os.getenv(
            "ROOVY_BASE_URL", "https://api.openai.com/v1"
        )
        self.model = model or os.getenv("ROOVY_MODEL", "gpt-4o")
        self._use_docs = use_docs
        self._custom_docs_path = custom_docs_path
        self._custom_docs = custom_docs

        if not self.api_key:
            raise ValueError(
                "API key is required. Pass api_key or set ROOVY_API_KEY environment variable."
            )

        self._client = OpenAI(base_url=self.base_url, api_key=self.api_key)

    def generate(
        self,
        prompt: str,
        *,
        max_retries: int = 3,
        temperature: float = 0.2,
        system_prompt: Optional[str] = None,
    ) -> WhatsAppFlow:
        """Generate a WhatsApp Flow from a natural language description.

        Args:
            prompt: Natural language description of the desired flow.
            max_retries: Number of retry attempts on validation failure.
            temperature: LLM temperature (0.0-1.0).
            system_prompt: Custom system prompt. Uses default if not provided.

        Returns:
            A validated WhatsAppFlow object.

        Raises:
            RuntimeError: If generation fails after all retries.
        """
        last_error: Optional[Exception] = None
        last_validation_errors: str = ""

        # Build system prompt with optional documentation
        if system_prompt:
            sys_prompt = system_prompt
        else:
            sys_prompt = get_enhanced_prompt(
                SYSTEM_PROMPT,
                use_bundled_docs=self._use_docs,
                custom_docs_path=self._custom_docs_path,
                custom_docs=self._custom_docs,
            )

        for attempt in range(1, max_retries + 1):
            try:
                messages: list[dict[str, str]] = [
                    {"role": "system", "content": sys_prompt},
                    {"role": "user", "content": prompt},
                ]

                if attempt > 1 and last_validation_errors:
                    messages.append(
                        {
                            "role": "user",
                            "content": (
                                "Previous attempt had validation errors. Fix these issues:\n"
                                f"{last_validation_errors}\n\nGenerate corrected JSON only."
                            ),
                        }
                    )

                response = self._client.chat.completions.create(
                    model=self.model,
                    messages=messages,  # type: ignore[arg-type]
                    temperature=temperature,
                    max_tokens=10_000,
                )

                content = (
                    response.choices[0].message.content if response.choices else None
                )
                if not content:
                    raise RuntimeError("Empty response from LLM")

                json_str = _clean_json_response(content)

                try:
                    parsed: Any = json.loads(json_str)
                except json.JSONDecodeError as parse_err:
                    raise RuntimeError(f"JSON parse error: {parse_err}") from parse_err

                validation = validate_flow(parsed)
                if validation["success"]:
                    return validation["data"]

                err: ValidationError = validation["error"]
                last_validation_errors = _format_validation_errors(err)
                raise RuntimeError(
                    f"Schema validation failed:\n{last_validation_errors}"
                )

            except Exception as exc:
                last_error = exc

        raise RuntimeError(
            f"Failed after {max_retries} attempts. Last error: {last_error}"
            if last_error
            else f"Failed after {max_retries} attempts."
        )

    def generate_stream(
        self,
        prompt: str,
        on_chunk: Optional[Callable[[str], None]] = None,
        *,
        temperature: float = 0.2,
        system_prompt: Optional[str] = None,
    ) -> WhatsAppFlow:
        """Generate a WhatsApp Flow with streaming output.

        Args:
            prompt: Natural language description of the desired flow.
            on_chunk: Optional callback called with each streamed chunk.
            temperature: LLM temperature (0.0-1.0).
            system_prompt: Custom system prompt. Uses default if not provided.

        Returns:
            A validated WhatsAppFlow object.

        Raises:
            RuntimeError: If generation or validation fails.
        """
        # Build system prompt with optional documentation
        if system_prompt:
            sys_prompt = system_prompt
        else:
            sys_prompt = get_enhanced_prompt(
                SYSTEM_PROMPT,
                use_bundled_docs=self._use_docs,
                custom_docs_path=self._custom_docs_path,
                custom_docs=self._custom_docs,
            )

        stream = self._client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": sys_prompt},
                {"role": "user", "content": prompt},
            ],  # type: ignore[arg-type]
            temperature=temperature,
            max_tokens=10_000,
            stream=True,
        )

        full_content = ""
        for chunk in stream:
            delta = chunk.choices[0].delta.content or ""
            full_content += delta
            if on_chunk is not None and delta:
                on_chunk(delta)

        json_str = _clean_json_response(full_content)
        parsed: Any = json.loads(json_str)

        validation = validate_flow(parsed)
        if not validation["success"]:
            err: ValidationError = validation["error"]
            raise RuntimeError(f"Validation failed: {_format_validation_errors(err)}")

        return validation["data"]

    def generate_template(
        self,
        prompt: str,
        *,
        max_retries: int = 3,
        temperature: float = 0.2,
        system_prompt: Optional[str] = None,
    ) -> TemplateResult:
        """Generate a WhatsApp Message Template from a natural language description.

        Returns both the template creation JSON (for registering with Meta)
        and the send message JSON (for sending to a user).

        Args:
            prompt: Natural language description of the desired template.
            max_retries: Number of retry attempts on validation failure.
            temperature: LLM temperature (0.0-1.0).
            system_prompt: Custom system prompt. Uses default if not provided.

        Returns:
            A validated TemplateResult with .creation and .send properties.

        Raises:
            RuntimeError: If generation fails after all retries.
        """
        last_error: Optional[Exception] = None
        last_validation_errors: str = ""

        if system_prompt:
            sys_prompt = system_prompt
        else:
            sys_prompt = TEMPLATE_SYSTEM_PROMPT + "\n\n" + TEMPLATE_DOCS

        for attempt in range(1, max_retries + 1):
            try:
                messages: list[dict[str, str]] = [
                    {"role": "system", "content": sys_prompt},
                    {"role": "user", "content": prompt},
                ]

                if attempt > 1 and last_validation_errors:
                    messages.append(
                        {
                            "role": "user",
                            "content": (
                                "Previous attempt had validation errors. Fix these issues:\n"
                                f"{last_validation_errors}\n\nGenerate corrected JSON only."
                            ),
                        }
                    )

                response = self._client.chat.completions.create(
                    model=self.model,
                    messages=messages,  # type: ignore[arg-type]
                    temperature=temperature,
                    max_tokens=10_000,
                )

                content = (
                    response.choices[0].message.content if response.choices else None
                )
                if not content:
                    raise RuntimeError("Empty response from LLM")

                json_str = _clean_json_response(content)

                try:
                    parsed: Any = json.loads(json_str)
                except json.JSONDecodeError as parse_err:
                    raise RuntimeError(f"JSON parse error: {parse_err}") from parse_err

                validation = validate_template_result(parsed)
                if validation["success"]:
                    return validation["data"]

                err = validation["error"]
                last_validation_errors = _format_validation_errors(err)
                raise RuntimeError(
                    f"Schema validation failed:\n{last_validation_errors}"
                )

            except Exception as exc:
                last_error = exc

        raise RuntimeError(
            f"Failed after {max_retries} attempts. Last error: {last_error}"
            if last_error
            else f"Failed after {max_retries} attempts."
        )

    def generate_template_stream(
        self,
        prompt: str,
        on_chunk: Optional[Callable[[str], None]] = None,
        *,
        temperature: float = 0.2,
        system_prompt: Optional[str] = None,
    ) -> TemplateResult:
        """Generate a WhatsApp Message Template with streaming output.

        Args:
            prompt: Natural language description of the desired template.
            on_chunk: Optional callback called with each streamed chunk.
            temperature: LLM temperature (0.0-1.0).
            system_prompt: Custom system prompt. Uses default if not provided.

        Returns:
            A validated TemplateResult with .creation and .send properties.

        Raises:
            RuntimeError: If generation or validation fails.
        """
        if system_prompt:
            sys_prompt = system_prompt
        else:
            sys_prompt = TEMPLATE_SYSTEM_PROMPT + "\n\n" + TEMPLATE_DOCS

        stream = self._client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": sys_prompt},
                {"role": "user", "content": prompt},
            ],  # type: ignore[arg-type]
            temperature=temperature,
            max_tokens=10_000,
            stream=True,
        )

        full_content = ""
        for chunk in stream:
            delta = chunk.choices[0].delta.content or ""
            full_content += delta
            if on_chunk is not None and delta:
                on_chunk(delta)

        json_str = _clean_json_response(full_content)
        parsed: Any = json.loads(json_str)

        validation = validate_template_result(parsed)
        if not validation["success"]:
            err = validation["error"]
            raise RuntimeError(f"Validation failed: {_format_validation_errors(err)}")

        return validation["data"]
