import shutil
import os
import streamlit as st
from huggingface_hub import snapshot_download

from src.config_loader import load_config, get_api_key
from src.ingest import get_chroma_collection, ingest_documents
from src.retriever import clear_collection_cache, retrieve
from src.llm import list_models
from src.peri import APPROVED_OPTIONS
from src.query_engine import understand_query, categorize_query, init_query, analyze_committment_to_vc
from src.prompts import dict_to_string

snapshot_download(repo_id="CGIAR/peri-kb", 
                  repo_type="dataset", 
                  allow_patterns="ug/*",
                  token=os.getenv('HF_TOKEN'), 
                  local_dir="./"
                  )
shutil.copytree("./ug", "./", dirs_exist_ok=True)


def render_sidebar():
    """Render sidebar with provider/model selectors, web search toggle, and KB stats."""
    cfg = st.session_state.cfg

    with st.sidebar:
        st.header("Settings")

        # --- Provider dropdown ---
        providers = ["openai", "anthropic", "gemini", "meta-llama"]
        current_provider = cfg.get("llm", {}).get("provider", "openai")
        provider_index = providers.index(current_provider) if current_provider in providers else 0

        provider = st.selectbox(
            "LLM Provider",
            providers,
            index=provider_index,
            key="sidebar_provider",
        )

        # Update cfg in session when provider changes
        if provider != cfg.get("llm", {}).get("provider"):
            cfg.setdefault("llm", {})["provider"] = provider

        # --- Model dropdown (cached per provider) ---
        # Invalidate model cache if provider changed
        prev_provider_key = "prev_provider"
        if st.session_state.get(prev_provider_key) != provider:
            for p in providers:
                st.session_state.pop(f"models_{p}", None)
            st.session_state[prev_provider_key] = provider

        models_cache_key = f"models_{provider}"
        if models_cache_key not in st.session_state:
            api_key = get_api_key(cfg, provider)
            if api_key:
                try:
                    st.session_state[models_cache_key] = list_models(provider, api_key)
                except Exception:
                    st.session_state[models_cache_key] = []
            else:
                st.session_state[models_cache_key] = []

        available_models = st.session_state[models_cache_key]
        current_model = cfg.get("llm", {}).get("model", "")

        if available_models:
            model_index = (
                available_models.index(current_model)
                if current_model in available_models
                else 0
            )
            model = st.selectbox(
                "Model",
                available_models,
                index=model_index,
                key="sidebar_model",
            )
        else:
            model = st.text_input(
                "Model",
                value=current_model,
                key="sidebar_model_text",
            )

        # Update cfg in session when model changes
        if model != cfg.get("llm", {}).get("model"):
            cfg.setdefault("llm", {})["model"] = model

        # --- Web search toggle ---
        web_enabled = cfg.get("web_search", {}).get("enabled", False)
        web_toggle = st.toggle("Web search", value=web_enabled, key="sidebar_web_search")
        cfg.setdefault("web_search", {})["enabled"] = web_toggle

        st.divider()

        # --- Knowledge base stats ---
        st.subheader("Knowledge Base")
        try:
            collection = get_chroma_collection(cfg)
            chunk_count = collection.count()
            st.metric("Chunks indexed", chunk_count)
        except Exception as e:
            st.warning(f"Could not read knowledge base: {e}")
            chunk_count = 0

        # --- Re-ingest button ---
        if st.button("Re-ingest documents", use_container_width=True):
            st.info("Ingestion may take a few minutes for large document collections...")
            with st.spinner("Ingesting documents..."):
                try:
                    count = ingest_documents(cfg)
                    clear_collection_cache()
                    st.success(f"Ingested {count} chunks.")
                    # Clear cached data so it refreshes after re-ingest
                    st.session_state.pop("kb_welcome_summary", None)
                    st.rerun()
                except Exception as e:
                    st.error(f"Ingestion failed: {e}")


def render_chat():
    """Render the chat interface with message history and input."""
    cfg = st.session_state.cfg

    # Display chat history
    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    # Chat input
    user_input = st.chat_input("Ask a question about your knowledge base...",
                            accept_file="multiple", 
                            file_type=["docx", "csv", "xlsx", "xls", "pdf", "rds", "rda", 
                                       "tsv", "sav", "dta", "txt", "md", "json", "do"]
                        )

    if user_input:
        prompt = user_input.text
        uploaded_files = user_input.files
        # Show and store user message
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        # ── Categorize user query ───────────────────
        save_dir = "uploaded_do_files"
        os.makedirs(save_dir, exist_ok=True)

        # Exclude the just-appended user message to avoid sending
        # the current question twice (once in history, once as query)
        cat_cfg = cfg.get("query_categorization", {})
        max_history = cat_cfg.get("max_history", 6)

        # ── Query understanding ─────────────────────────────────────────
        qu_cfg = cfg.get("query_understanding", {})
        qu_enabled = qu_cfg.get("enabled", True)
        max_history = qu_cfg.get("max_history", 6)

        # ── Check if this is a clarification response ───────────────────
        unresolved = st.session_state.unresolved_category
        if unresolved is not None:
            # This prompt is the user's clarification answer
            # Include the clarification question for context
            unresolved_question = st.session_state.get("unresolved_category_question", "")
            if unresolved_question:
                combined_cat = f"{unresolved} (Clarification: Q: {unresolved_question} A: {prompt})"
            else:
                combined_cat = f"{unresolved} — {prompt}"
            st.session_state.unresolved_category_question = None
            original_query_cat = unresolved
            st.session_state.unresolved_category = None
        else:
            combined_cat = prompt
            original_query_cat = prompt
            st.session_state.resolution_rounds = 0
            unresolved_question = []

        prior_messages = st.session_state.messages[:-1]
        history = [
            {"role": m["role"], "content": m["content"]}
            for m in prior_messages[-max_history:]
        ]
        try:
            qinit_result = init_query(combined_cat, cfg, history)
            print("result: ", qinit_result)
            if qinit_result.get("country", None) is None:
                
                with st.chat_message("assistant"):
                    st.markdown("Please specify the country for which you'd like to run a PERI analysis")
                return
            if qinit_result.get("value_chain", None) is None and qinit_result.get("investment", None) is None:
                
                with st.chat_message("assistant"):
                    st.markdown(f"Please specify the value chain or investment area in {qinit_result.get('country', None)} for which you'd like to run a PERI analysis")
                return
            qcat_result = categorize_query(combined_cat, cfg, history)
            print("result: ", qcat_result)
        except Exception as e:
            print(e)
            qinit_result = {"country": None, "value_chain": None, "investment": None}
            qcat_result = {"category": "pillar_1", "action": "unresolved"}

        if not isinstance(qinit_result.get("country", None), list) and qinit_result.get("country", None)==None:
            resolution_msg = f"It seems there is no country specified for the analysis. Currently the PERI framework supports analysis for the countries listed below:\n\n  {', '.join([c.capitalize() for c in APPROVED_OPTIONS.get('countries')])}\n\n We are also continuously \
                            working to expand the framework and you can submit a form if the country you would like to run the analysis on is not included. In the meantime please let me know if you would like to run the anlysis for one of the included countries."
                        
        if qinit_result.get("country", None)[0].lower() not in [c.lower() for c in APPROVED_OPTIONS.get('countries')]:
            resolution_msg = f"It seems you are trying to run a PERI analysis for {qinit_result.get('country', None)[0].capitalize()}! Currently the PERI framework only supports analysis for the countries listed below:\n\n  {', '.join([c.capitalize() for c in APPROVED_OPTIONS.get('countries')])}\n\n We are continuously \
                working to expand the framework and you can submit a form to request. In the meantime please let me know if you would like to run the anlysis for one of the included countries."

        if qinit_result.get("value_chain", None)[0]==None and qinit_result.get('investment', None)[0]==None:
            resolution_msg = f"It seems there is no value chain or investment area specified for the analysis. Please select one of the value chains or investment areas included in the current PERI framework."
        
        if qinit_result.get("value_chain", None)[0].lower() not in [c.lower() for v in APPROVED_OPTIONS.get('value_chains').values() for c in v]:
            resolution_msg = f"It seems you are trying to run a PERI analysis for {qinit_result.get('value_chain', None)[0]} in {qinit_result.get('country', None)[0].capitalize()}! Currently the PERI framework only supports analysis for the value chains listed below:\n\n  {dict_to_string(APPROVED_OPTIONS.get('value_chains'), 2)}\n\n We are continuously \
                working to expand the framework and you can submit a form to request. In the meantime please let me know if you would like to run the anlysis for one of the included countries."
                
        if qinit_result.get("investment", None)[0].lower() not in [i.lower() for i in APPROVED_OPTIONS.get('investments')]:
            resolution_msg = f"It seems you are trying to run a PERI analysis for {qinit_result.get('investment', None)[0]} investemnt in {qinit_result.get('country', None)[0].capitalize()}! Currently the PERI framework only supports analysis for the investment areas listed below:\n\n  {', '.join([c.capitalize() for c in APPROVED_OPTIONS.get('investments')])}\n\n We are continuously \
                working to expand the framework and you can submit a form to request. In the meantime please let me know if you would like to run the anlysis for one of the included countries."

        if qinit_result.get("country", None)[0] in APPROVED_OPTIONS.get('countries') and (qinit_result.get('value_chain', None)[0] in APPROVED_OPTIONS.get('value_chains') or qinit_result.get('investment', None)[0] in APPROVED_OPTIONS.get('investments')):
    
            pillar_dict = {
                "pillar_1":f"would like to understand whether {','.join(qinit_result.get('value_chain'))} aligns with the government's political incentives.",
                "pillar_2":f"would like to understand to what degree decisions are impacted by the lobbying of particular groups or by elite influence",
                "pillar_3":f"would like to understand if {','.join(qinit_result.get('value_chain'))} and/or investing in {','.join(qinit_result.get('investment'))} can be feasibly implemented given the broader institutional and policy environment",
                
            }
            resolution_msg = f"**Before I search, could you clarify?** Please let me know if you {' and '.join([pillar_dict[c] for c in qcat_result.get('category')])}."
                                 
        # st.session_state.unresolved_category_question = f"Please let me know if you {' and '.join([pillar_dict[c] for c in qcat_result.get('category')])}."
        st.session_state.unresolved_category_question = resolution_msg
        st.session_state.unresolved_category = original_query_cat
        st.session_state.resolution_rounds += 1
        st.session_state.messages.append({"role": "assistant", "content": resolution_msg})
        
        with st.chat_message("assistant"):
            st.markdown(resolution_msg)
        return