"""
loki.py

Databricks Model Serving (OpenAI-compatible) wrapper with:
- Loki.ask(): stateless call
- TradingChatSession: stateful commodity trading analytics chat
- SqlChatSession: stateful Databricks SQL generator chat

Important constraint:
- Gemini models only support ONE system prompt.
  => We must NOT send multiple system messages.
  => We fold summary + context blocks into a single system string.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

from ..workspaces.workspace import WorkspaceService

try:
    from openai import OpenAI as _OpenAI  # noqa: F401
except ImportError:
    _OpenAI = None  # type: ignore


def make_ai_client(api_key: str, base_url: str):
    """Late import so module can load even if openai isn't installed."""
    from openai import OpenAI
    return OpenAI(api_key=api_key, base_url=base_url)


__all__ = ["Loki", "TradingChatSession", "SqlChatSession"]


DEFAULT_TRADING_SYSTEM_PROMPT = """You are Loki: a conversational commodity trading analytics copilot.

Scope:
- Commodity trading analytics: curves, forwards, spreads, basis, hedging, risk, PnL explain, inventory, scheduling.
- Databricks-first workflows: Spark/Delta/Unity Catalog, Databricks SQL, performant Python.

Rules:
- Do NOT invent prices, positions, PnL, risk, or market facts not provided.
- State assumptions explicitly (units, time conventions, delivery months, calendars).
- Prefer actionable output (SQL + efficient Python). Avoid slow patterns.
- If data is missing, list exactly what you need and proceed with a reasonable template.

Style:
- Be concise, practical, performance-focused.
"""

DEFAULT_SQL_SYSTEM_PROMPT = """You are LokiSQL: a Databricks SQL generator for commodity trading analytics.

Hard rules:
- Output ONLY SQL unless the user explicitly asks for explanation.
- Use Databricks SQL / Spark SQL dialect.
- Prefer readable CTEs, explicit column lists, deterministic joins.
- Do NOT invent table/column names. If missing, use placeholders like <table>, <col>.
- If ambiguous, output best-effort SQL template with SQL comments (-- TODO ...) and placeholders.
- Performance: push filters early, avoid exploding joins, avoid SELECT *.

Default assumptions:
- Dates UTC unless specified.
"""


@dataclass
class Loki(WorkspaceService):
    """
    Loki wraps an OpenAI-compatible client pointing at Databricks Model Serving endpoints.
    """

    model: str = "databricks-gemini-2-5-flash"
    _ai_client: Optional[Any] = field(repr=False, hash=False, default=None)

    @property
    def ai_client(self):
        if self._ai_client is None:
            self._ai_client = self.make_aiclient()
        return self._ai_client

    def make_aiclient(self):
        host = self.workspace.host.rstrip("/")
        return make_ai_client(
            api_key=self.workspace.current_token(),
            base_url=f"{host}/serving-endpoints",
        )

    def ask(
        self,
        command: str,
        *,
        system: Optional[str] = None,
        max_tokens: int = 5000,
        temperature: Optional[float] = None,
        extra_messages: Optional[List[Dict[str, str]]] = None,
        **kwargs: Any,
    ) -> str:
        """
        Stateless single call to the model.

        NOTE (Gemini constraint):
        - Provide at most ONE system prompt (i.e., a single system message).
        - Do not pass additional messages with role="system".
        """
        messages: List[Dict[str, str]] = []
        if system:
            messages.append({"role": "system", "content": system})

        if extra_messages:
            # IMPORTANT: caller must not include any "system" roles here for Gemini models
            messages.extend(extra_messages)

        messages.append({"role": "user", "content": command})

        params: Dict[str, Any] = dict(
            model=self.model,
            messages=messages,
            max_tokens=max_tokens,
            **kwargs,
        )
        if temperature is not None:
            params["temperature"] = temperature

        resp = self.ai_client.chat.completions.create(**params)
        return resp.choices[0].message.content or ""

    def new_trading_chat(
        self,
        *,
        system_prompt: str = DEFAULT_TRADING_SYSTEM_PROMPT,
        max_context_turns: int = 20,
        max_context_chars: int = 120_000,
    ) -> "TradingChatSession":
        return TradingChatSession(
            loki=self,
            system_prompt=system_prompt,
            max_context_turns=max_context_turns,
            max_context_chars=max_context_chars,
        )

    def new_sql_chat(
        self,
        *,
        system_prompt: str = DEFAULT_SQL_SYSTEM_PROMPT,
        max_context_turns: int = 20,
        max_context_chars: int = 120_000,
    ) -> "SqlChatSession":
        return SqlChatSession(
            loki=self,
            system_prompt=system_prompt,
            max_context_turns=max_context_turns,
            max_context_chars=max_context_chars,
        )


@dataclass
class _BaseChatSession:
    """
    Stateful session that maintains history + injected context blocks.

    Gemini constraint:
    - We must fold ALL system content into one system string.
    - Therefore summary/context_blocks are concatenated into the system prompt.
    """
    loki: Loki
    system_prompt: str

    history: List[Dict[str, str]] = field(default_factory=list)
    summary: Optional[str] = None
    context_blocks: List[str] = field(default_factory=list)

    max_context_turns: int = 20
    max_context_chars: int = 120_000

    def reset(self) -> None:
        self.history.clear()
        self.summary = None
        self.context_blocks.clear()

    def add_context(self, title: str, payload: Union[str, Dict[str, Any], List[Any]]) -> None:
        if isinstance(payload, str):
            payload_str = payload
        else:
            payload_str = json.dumps(payload, ensure_ascii=False, indent=2)

        self.context_blocks.append(f"[Context: {title}]\n{payload_str}".strip())
        self._trim()

    def _estimate_chars(self, msgs: List[Dict[str, str]]) -> int:
        return sum(len(m.get("content", "")) for m in msgs)

    def _build_system(self, extra_system: Optional[str] = None) -> str:
        parts: List[str] = [self.system_prompt.strip()]
        if extra_system:
            parts.append(extra_system.strip())
        if self.summary:
            parts.append(f"[ConversationSummary]\n{self.summary}".strip())
        if self.context_blocks:
            parts.append("\n\n".join(self.context_blocks).strip())
        return "\n\n".join(p for p in parts if p)

    def _trim(self) -> None:
        # Turn trim (keep last N turns => N*2 messages)
        if self.max_context_turns > 0:
            max_msgs = self.max_context_turns * 2
            if len(self.history) > max_msgs:
                self.history = self.history[-max_msgs:]

        # Char trim: shrink history first, then context blocks if needed
        def total_chars() -> int:
            sys_len = len(self._build_system())
            return sys_len + self._estimate_chars(self.history)

        while total_chars() > self.max_context_chars and self.history:
            self.history = self.history[1:]

        while total_chars() > self.max_context_chars and self.context_blocks:
            self.context_blocks = self.context_blocks[1:]


@dataclass
class TradingChatSession(_BaseChatSession):
    """
    Commodity trading analytics chat session.
    Optionally returns structured JSON for downstream automation.
    """

    def chat(
        self,
        user_text: str,
        *,
        structured: bool = True,
        max_tokens: int = 12000,
        temperature: Optional[float] = None,
        **kwargs: Any,
    ) -> Union[str, Dict[str, Any]]:
        self._trim()

        extra_system = None
        if structured:
            extra_system = (
                "Respond ONLY as valid JSON with keys: "
                "final_answer (string), assumptions (array of strings), data_needed (array of strings), "
                "sql (string or null), python (string or null). "
                "No markdown. No extra keys."
            )

        system = self._build_system(extra_system=extra_system)

        assistant_text = self.loki.ask(
            user_text,
            system=system,
            extra_messages=self.history,  # NOTE: history must contain no system roles
            max_tokens=max_tokens,
            temperature=temperature,
            **kwargs,
        )

        self.history.append({"role": "user", "content": user_text})
        self.history.append({"role": "assistant", "content": assistant_text})
        self._trim()

        if structured:
            parsed = _try_parse_json_object(assistant_text)
            if parsed is not None:
                return parsed

        return assistant_text


@dataclass
class SqlChatSession(_BaseChatSession):
    """
    SQL-only conversational session that generates Databricks SQL.

    Uses a single system message with strict instructions to output SQL only.
    """

    def generate_sql(
        self,
        request: str,
        *,
        max_tokens: int = 12000,
        temperature: Optional[float] = None,
        sql_only: bool = True,
        **kwargs: Any,
    ) -> str:
        self._trim()

        extra_system = None
        if sql_only:
            extra_system = (
                "Reminder: Output ONLY SQL. "
                "If ambiguity exists, use SQL comments (-- TODO ...) and placeholders, but still output SQL only."
            )

        system = self._build_system(extra_system=extra_system)

        sql = self.loki.ask(
            request,
            system=system,
            extra_messages=self.history,  # history must contain no system roles
            max_tokens=max_tokens,
            temperature=temperature,
            **kwargs,
        ).strip()

        self.history.append({"role": "user", "content": request})
        self.history.append({"role": "assistant", "content": sql})
        self._trim()

        return _strip_sql_fences(sql)


def _strip_sql_fences(text: str) -> str:
    t = text.strip()
    if t.startswith("```"):
        lines = t.splitlines()
        lines = lines[1:]  # drop ``` or ```sql
        if lines and lines[-1].strip().startswith("```"):
            lines = lines[:-1]
        return "\n".join(lines).strip()
    return t


def _strip_markdown_fences(text: str) -> str:
    """
    Remove ```lang ... ``` fences if present.
    Keeps inner content unchanged.
    """
    t = text.strip()
    if not t.startswith("```"):
        return t

    lines = t.splitlines()
    if not lines:
        return t

    # Drop first line: ``` or ```json
    lines = lines[1:]

    # Drop last line if it's ```
    if lines and lines[-1].strip().startswith("```"):
        lines = lines[:-1]

    return "\n".join(lines).strip()


def _try_parse_json_object(text: str) -> Optional[Dict[str, Any]]:
    t = _strip_markdown_fences(text).strip()

    # Best effort extraction if there's extra junk around JSON
    if not t.startswith("{"):
        start = t.find("{")
        end = t.rfind("}")
        if start != -1 and end != -1 and end > start:
            t = t[start : end + 1]

    try:
        obj = json.loads(t)
    except Exception:
        return None

    if not isinstance(obj, dict):
        return None

    required = {"final_answer", "assumptions", "data_needed", "sql", "python"}
    if not required.issubset(set(obj.keys())):
        return None

    return obj
