from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol

import psycopg

from dotenv import load_dotenv

from rag_agent.config import AppConfig
from rag_agent.index.embeddings import EmbeddingClient
from rag_agent.retrieval.search import search_similar

_repo_root = Path(__file__).resolve().parent.parent.parent
load_dotenv(_repo_root / ".env")


class LLMClient(Protocol):
    def generate(self, prompt: str, model: str) -> str:
        raise NotImplementedError


@dataclass
class StubLLMClient:
    def generate(self, prompt: str, model: str) -> str:
        return (
            "LLM client is not configured. "
            "Set GIGACHAT_CREDENTIALS in .env for GigaChat answers."
        )


class GigaChatLLMClient:
    """LLM generation via GigaChat API. Credentials from env GIGACHAT_CREDENTIALS."""

    def __init__(
        self,
        credentials: str,
        model: str = "GigaChat",
        verify_ssl_certs: bool = False,
    ) -> None:
        self._credentials = credentials.strip()
        self._model = model
        self._verify_ssl_certs = verify_ssl_certs

    def generate(self, prompt: str, model: str) -> str:
        from gigachat import GigaChat

        use_model = model or self._model
        with GigaChat(
            credentials=self._credentials,
            model=use_model,
            verify_ssl_certs=self._verify_ssl_certs,
        ) as giga:
            response = giga.chat(prompt)
        return (response.choices[0].message.content or "").strip()


def get_llm_client(config: AppConfig) -> LLMClient:
    """Return GigaChat LLM client if credentials set, else stub."""
    credentials = os.getenv("GIGACHAT_CREDENTIALS", "").strip()
    if credentials:
        return GigaChatLLMClient(
            credentials=credentials,
            model=config.llm_model,
            verify_ssl_certs=os.getenv("GIGACHAT_VERIFY_SSL", "false").lower()
            in ("1", "true", "yes"),
        )
    return StubLLMClient()


def build_prompt(question: str, contexts: list[str]) -> str:
    joined = "\n\n".join(contexts)
    return (
        "You are a RAG assistant. Use the context below to answer the question.\n\n"
        f"Context:\n{joined}\n\n"
        f"Question: {question}\nAnswer:"
    )


def answer_query(
    conn: psycopg.Connection,
    config: AppConfig,
    embedding_client: EmbeddingClient,
    llm_client: LLMClient,
    question: str,
    top_k: int = 5,
    story_id: int | None = None,
) -> str:
    answer, _ = answer_query_with_stats(
        conn, config, embedding_client, llm_client, question, top_k, story_id
    )
    return answer


def answer_query_with_stats(
    conn: psycopg.Connection,
    config: AppConfig,
    embedding_client: EmbeddingClient,
    llm_client: LLMClient,
    question: str,
    top_k: int = 5,
    story_id: int | None = None,
) -> tuple[str, dict]:
    """Like answer_query but returns (answer, stats) for logging. stats: query_embeddings, chunks_found, answer."""
    query_embeddings = embedding_client.embed_texts([question])
    results = search_similar(
        conn, query_embeddings[0], top_k=top_k, story_id=story_id
    )
    contexts = [f"Source: {item.path}\n{item.content}" for item in results]
    prompt = build_prompt(question, contexts)
    answer = llm_client.generate(prompt, model=config.llm_model)
    stats = {
        "query_embeddings": len(query_embeddings),
        "chunks_found": len(results),
        "answer": answer,
    }
    return answer, stats
