from pydantic import BaseModel, Field
from typing import Literal
from core.llm_router import get_llm
from langchain_core.prompts import PromptTemplate
from rag_pipeline import get_hybrid_retriever, rerank_documents
import logging
from tenacity import retry, stop_after_attempt, wait_exponential

logger = logging.getLogger(__name__)


class ExpenseEvaluationResponse(BaseModel):
    czy_wydatek_kwalifikowalny: bool = Field(
        description="Zwróć True jeśli wydatek jest w 100% zgodny z regulaminem i wytycznymi programu (kwalifikowalny)."
    )
    uzasadnienie_prawne: str = Field(
        description="Cytat lub konkretne odwołanie do regulaminu uzasadniające kwalifikowalność lub jej brak."
    )
    kategoria_badan: Literal[
        "badania przemysłowe",
        "prace rozwojowe",
        "prace przedwdrożeniowe",
        "brak/nie dotyczy",
    ] = Field(
        description="Wybierz do jakiej kategorii zgodnie z polskim/unijnym prawem należy ten wydatek. Wybierz 'brak/nie dotyczy' tylko jeśli wydatek jest całkowicie poza B+R."
    )
    intensywnosc_pomocy: float = Field(
        description="Zwróć w formie wartości zmiennoprzecinkowej np. 0.50 (co oznacza 50%), 0.80 (co oznacza 80%) bazując na wielkości firmy i rodzaju badań. 0.0 oznacza wydatek niekwalifikowalny."
    )


def evaluate_project_expense(
    expense_description: str,
    expense_amount: float,
    project_title: str,
    program_name: str,
    company_size: str,
    tenant_id: str = None,
) -> ExpenseEvaluationResponse:
    """
    Agent ds. Oceny Kwalifikowalności (FAZA 4).
    Wymusza twarde, ustrukturyzowane ramy JSON za pomocą Pydantic.
    Opiera się na wiedzy RAG dotyczącej wybranego programu.
    """

    # Próba załadowania kontekstu z RAG - Hard Filtering na aktualną perspektywę
    # Domyślnie wyszukujemy tylko w najnowszej perspektywie (FAZA 3, zapobieganie aplikacji starych przepisów)
    hard_filter = {"rok_perspektywy": {"$eq": "2021-2027"}}
    if program_name:
        # Operator $and dla Pinecone Vector Store
        hard_filter = {
            "$and": [
                {"program_name": {"$eq": program_name}},
                {"rok_perspektywy": {"$eq": "2021-2027"}},
            ]
        }

    context_text = "Brak specyficznego regulaminu programu w bazie."
    try:
        retriever = get_hybrid_retriever(
            k=10, metadata_filter=hard_filter, namespace=tenant_id
        )
        if retriever:
            query_for_rag = f"kwalifikowalność wydatku badania kategoria intensywność dotacji pomoc publiczna: {expense_description}"
            docs = retriever.invoke(query_for_rag)
            reranked_docs = rerank_documents(query_for_rag, docs, top_n=4)
            context_text = "\n\n".join(
                [
                    f"[ŹRÓDŁO: {d.metadata.get('source', 'Brak')}]: {d.page_content}"
                    for d in reranked_docs
                ]
            )
    except Exception as e:
        logger.error(f"[ExpenseEvaluator] Error fetching RAG context: {str(e)}")

    template = """
    Jesteś Głównym Prawnikiem i Audytorem Dotacyjnym oceniającym kwalifikowalność wydatków.
    Oceniasz pojedynczy wydatek dla projektu w ramach programu: {program_name}.
    Wielkość przedsiębiorstwa wnioskodawcy: {company_size}.
    
    Opis wydatku do weryfikacji:
    "{expense_description}" (Kwota: {expense_amount} PLN)
    
    Kontekst z regulaminów z bazy wiedzy:
    --------------------------------------------------
    {context}
    --------------------------------------------------
    
    Zasady:
    1. Przeanalizuj czy podany wydatek kwalifikuje się do objęcia wsparciem zgodnie z bazą wiedzy.
    2. Określ kategorię badań dla wydatku, zgodnie z definicjami (badania przemysłowe, prace rozwojowe, przedwdrożeniowe).
    3. Jeśli wydatek jest kwalifikowalny, przypisz prawidłową intensywność pomocy (zazwyczaj mniejszy procent dla prac rozwojowych/dużych firm, większy dla badań przemysłowych/MŚP).
    4. Podaj bardzo precyzyjne uzasadnienie prawne odnoszące się do regulaminu.
    """

    prompt = PromptTemplate.from_template(template)

    # LLM z typowaniem - GPT-4o jest dużo lepszy do takich zadań analitycznych
    structured_llm = get_llm(
        task_type="legal_audit", structured_output_schema=ExpenseEvaluationResponse
    )

    chain = prompt | structured_llm

    @retry(
        stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10)
    )
    def _invoke_chain():
        return chain.invoke(
            {
                "program_name": program_name or "Ogólne zasady dotacji B+R",
                "company_size": company_size or "MŚP (nieokreślona wielkość)",
                "expense_description": expense_description,
                "expense_amount": expense_amount,
                "context": context_text,
            }
        )

    result = _invoke_chain()

    return result