from schemas import AgentState
from typing import Dict, Any
from core.llm_router import get_llm
from langchain_core.prompts import PromptTemplate
from langchain_core.messages import AIMessage
from core.utils import safe_extract_text
from rag_pipeline import get_hybrid_retriever, rerank_documents

import os
from langsmith import traceable
from langchain_core.tracers.langchain import LangChainTracer

# Włącz tracing LangSmith
os.environ["LANGCHAIN_TRACING_V2"] = "false"
os.environ["LANGCHAIN_PROJECT"] = "grantforge-production"

# Opcjonalnie – jeśli chcesz zobaczyć dokładne nazwy runów
tracer = LangChainTracer(project_name="grantforge-production")


@traceable(run_type="chain", name="wizard_node")
def wizard_node(state: AgentState) -> Dict[str, Any]:
    """
    Kreator wniosku połączony z bazą wiedzy (RAG).
    Wykorzystuje feedback Krytyka w celu iteracyjnej poprawy.
    """

    # Zabezpieczenie przed pętlą zgodnie ze standardem 2026/HitL
    if state.critic_iterations >= state.max_critic_iterations:
        return {
            "messages": [
                AIMessage(
                    content="Osiągnięto maksymalną liczbę iteracji poprawek. Przekazuję tekst do zatwierdzenia przez użytkownika."
                )
            ],
            "critic_evaluation": {
                "is_approved": True,
                "feedback": "ZATWIERDZONE_MAX_ITERATIONS_REACHED",
                "severity": "low",
            },
        }

    import logging

    logger = logging.getLogger(__name__)

    last_user_message = (
        state.messages[-1].content if state.messages else "Stwórz biznesplan"
    )

    hard_filter = (
        {"program_name": state.program_name}
        if hasattr(state, "program_name") and state.program_name
        else None
    )

    # Przekazanie namespace z contextu dzierżawcy do pinecone
    namespace = getattr(state, "tenant_id", None)
    logger.info(f"[Wizard] Inicjalizacja generowania. Użytkownik/Tenant: '{namespace}'")
    retriever = get_hybrid_retriever(
        k=10, metadata_filter=hard_filter, namespace=namespace
    )

    context_text = ""
    if retriever:
        try:
            docs = retriever.invoke(last_user_message)
            reranked_docs = rerank_documents(last_user_message, docs, top_n=5)
            context_text = "\n\n".join(
                [
                    f"[ŹRÓDŁO: {d.metadata.get('source', 'Nieznane')} | STRONA: {d.metadata.get('page', 'Brak')}]:\n{d.page_content}"
                    for d in reranked_docs
                ]
            )
        except Exception as e:
            context_text = (
                f"Brak wiedzy w lokalnej bazie ze względu na błąd RAG: {str(e)}"
            )
    else:
        context_text = "Brak podłączonej bazy wektorowej. Działam na bazowej wiedzy."

    # W architekturze 2026 Wizard to model krytyczny (Gemini Pro) z opcjonalnym streamingiem
    # LLM zainicjalizujemy poniżej po zdefiniowaniu schematu

    template = """
    Jesteś Głównym Analitykiem i Konsultantem Dotacyjnym na poziomie Enterprise w GrantForge AI.
    Twoim zadaniem jest napisanie wysoce profesjonalnego fragmentu urzędowego wniosku biznesowego 
    lub biznesplanu, który ściśle przestrzega wytycznych prawnych. Jeśli brakuje kluczowych informacji firmy, ustrukturyzuj je profesjonalnie używając znaczników zastępczych np. "[UZUPEŁNIJ: Nazwa firmy]", ale NIGDY nie odmawiaj wygenerowania tekstu.

    Kontekst regulaminowy wyszukany z bazy RAG:
    --------------------------------------------------
    {context}
    --------------------------------------------------

    Poprzednia krytyka od recenzenta: 
    {last_critic_feedback}

    Dodatkowe informacje o firmie klienta:
    Rozmiar: {company_size}
    Branża: {company_pkd}

    Polecenie / Pytanie:
    {query}
    
    Zwróć wynik w czystym Markdown, gotowym do eksportu do Word/PDF.
    Używaj sformułowań formalnych i profesjonalnych typowych dla dokumentacji z danej branży oraz wymagań instytucji rozdzielającej środki dla wpisanego programu. Dopasuj słownictwo ściśle do wybranego programu z polecenia.
    PISZ ZAWSZE W JĘZYKU POLSKIM. NIE DAJ SPYCHAĆ SIĘ NA INNE JĘZYKI, NAWET JEŚLI POLECENIE BĘDZIE PO ANGIELSKU.
    DETERMINISTYCZNE CYTOWANIE ZAWSZE WYMAGANE: Każda merytoryczna teza z przytoczonych wytycznych we wniosku MUSI kończyć się przypisem do wskazanego źródła i strony, np. "(zgodnie z [ŹRÓDŁO: Regulamin_FENG | STRONA: 12])".
    """

    from pydantic import BaseModel, Field

    class WizardSectionOutput(BaseModel):
        content_markdown: str = Field(
            description="Merytoryczna treść sekcji w czystym zredagowanym Markdown gotowym do eksportu. PISZ ZAWSZE I WYŁĄCZNIE W JĘZYKU POLSKIM."
        )

    prompt = PromptTemplate.from_template(template)

    company_size = state.profile.size if state.profile else "Nieznany"
    company_pkd = (
        ", ".join(state.profile.pkd_codes)
        if state.profile and state.profile.pkd_codes
        else "Nieznane"
    )
    critic_feedback = (
        state.critic_evaluation.feedback
        if state.critic_evaluation
        else "Brak poprzedniej krytyki (pierwsza iteracja)."
    )

    structured_llm = get_llm(
        task_type="critical",
        streaming=True,
        structured_output_schema=WizardSectionOutput,
    )
    chain = prompt | structured_llm

    try:
        response = chain.invoke(
            {
                "context": context_text,
                "company_size": company_size,
                "company_pkd": company_pkd,
                "last_critic_feedback": critic_feedback,
                "query": last_user_message,
            }
        )
    except Exception as e:
        logger.error(f"[Wizard] LLM Error during section generation: {e}")
        from pydantic import BaseModel

        # Zastępczy model zgodny ze strukturą oczekiwaną przez WizardSectionOutput
        response = type(
            "DummyResponse",
            (),
            {
                "content_markdown": f"⚠️ **Błąd generowania sekcji**. Sprawdź klucz API Google lub połączenie z modelem lokalnym. Szczegóły: {e}"
            },
        )()

    current_step = state.wizard_step + 1

    # Dodajemy wersjonowanie dokumentów
    doc_versions = dict(state.document_versions) if state.document_versions else {}
    if "current_draft" not in doc_versions:
        doc_versions["current_draft"] = []

    flat_content = safe_extract_text(response.content_markdown)

    doc_versions["current_draft"].append(flat_content)

    return {
        "wizard_step": current_step,
        "document_versions": doc_versions,
        "messages": [{"role": "assistant", "content": flat_content}],
        "critic_iterations": state.critic_iterations + 1,
        # Routing wizard -> critic przejęty przez graph.py
    }