# interface.py

import streamlit as st
import logging
from typing import Dict

from langchain_core.messages import HumanMessage
from workflow import ResearchWorkflow
from config import ResearchConfig
from langchain_core.messages import AIMessage

logger = logging.getLogger(__name__)

class ResearchInterface:
    """
    Provides the Streamlit-based interface for executing the research workflow.
    """
    def __init__(self) -> None:
        self.workflow = ResearchWorkflow()
        self._initialize_interface()

    def _initialize_interface(self) -> None:
        st.set_page_config(
            page_title="NeuroResearch AI",
            layout="wide",
            initial_sidebar_state="expanded"
        )
        self._inject_styles()
        self._build_sidebar()
        self._build_main_interface()

    def _inject_styles(self) -> None:
        st.markdown(
            """
            <style>
            :root {
                --primary: #2ecc71;
                --secondary: #3498db;
                --background: #0a0a0a;
                --text: #ecf0f1;
            }
            .stApp {
                background: var(--background);
                color: var(--text);
                font-family: 'Roboto', sans-serif;
            }
            .stTextArea textarea {
                background: #1a1a1a !important;
                color: var(--text) !important;
                border: 2px solid var(--secondary);
                border-radius: 8px;
                padding: 1rem;
            }
            .stButton>button {
                background: linear-gradient(135deg, var(--primary), var(--secondary));
                border: none;
                border-radius: 8px;
                padding: 1rem 2rem;
                transition: all 0.3s;
            }
            .stButton>button:hover {
                transform: translateY(-2px);
                box-shadow: 0 4px 12px rgba(46, 204, 113, 0.3);
            }
            .stExpander {
                background: #1a1a1a;
                border: 1px solid #2a2a2a;
                border-radius: 8px;
                margin: 1rem 0;
            }
            </style>
            """,
            unsafe_allow_html=True
        )

    def _build_sidebar(self) -> None:
        with st.sidebar:
            st.title("🔍 Research Database")
            st.subheader("Technical Papers")
            for title, short in ResearchConfig.DOCUMENT_MAP.items():
                with st.expander(short):
                    st.markdown(f"```\n{title}\n```")
            st.subheader("Analysis Metrics")
            st.metric("Vector Collections", 2)
            st.metric("Embedding Dimensions", ResearchConfig.EMBEDDING_DIMENSIONS)
            with st.sidebar.expander("Collaboration Hub"):
                st.subheader("Live Research Team")
                st.write("👩💻 Researcher A")
                st.write("👨🔬 Researcher B")
                st.write("🤖 AI Assistant")
                st.subheader("Knowledge Graph")
                if st.button("🕸 View Current Graph"):
                    self._display_knowledge_graph()

    def _build_main_interface(self) -> None:
        st.title("🧠 NeuroResearch AI")
        query = st.text_area("Research Query:", height=200, placeholder="Enter technical research question...")
        domain = st.selectbox(
            "Select Research Domain:",
            options=[
                "Biomedical Research",
                "Legal Research",
                "Environmental and Energy Studies",
                "Competitive Programming and Theoretical Computer Science",
                "Social Sciences"
            ],
            index=0
        )
        if st.button("Execute Analysis", type="primary"):
            self._execute_analysis(query, domain)

    def _execute_analysis(self, query: str, domain: str) -> None:
        try:
            with st.spinner("Initializing Quantum Analysis..."):
                results = self.workflow.app.stream(
                    {
                        "messages": [HumanMessage(content=query)],
                        "context": {"domain": domain},
                        "metadata": {}
                    },
                    {"recursion_limit": 100}
                )
                for event in results:
                    self._render_event(event)
                st.success("✅ Analysis Completed Successfully")
        except Exception as e:
            st.error(
                f"""**Analysis Failed**  
{str(e)}  
Potential issues:
- Complex query structure
- Document correlation failure
- Temporal processing constraints"""
            )

    def _render_event(self, event: Dict) -> None:
        if 'ingest' in event:
            with st.container():
                st.success("✅ Query Ingested")
        elif 'retrieve' in event:
            with st.container():
                docs = event['retrieve']['context'].get('documents', [])
                st.info(f"📚 Retrieved {len(docs)} documents")
                with st.expander("View Retrieved Documents", expanded=False):
                    for idx, doc in enumerate(docs, start=1):
                        st.markdown(f"**Document {idx}**")
                        st.code(doc.page_content, language='text')
        elif 'analyze' in event:
            with st.container():
                content = event['analyze']['messages'][0].content
                with st.expander("Technical Analysis Report", expanded=True):
                    st.markdown(content)
        elif 'validate' in event:
            with st.container():
                content = event['validate']['messages'][0].content
                if "VALID" in content:
                    st.success("✅ Validation Passed")
                    with st.expander("View Validated Analysis", expanded=True):
                        st.markdown(content.split("Validation:")[0])
                else:
                    st.warning("⚠️ Validation Issues Detected")
                    with st.expander("View Validation Details", expanded=True):
                        st.markdown(content)
        elif 'enhance' in event:
            with st.container():
                content = event['enhance']['messages'][0].content
                with st.expander("Enhanced Multi-Modal Analysis Report", expanded=True):
                    st.markdown(content)

    def _display_knowledge_graph(self) -> None:
        st.write("Knowledge Graph visualization is not implemented yet.")

class ResearchInterfaceExtended(ResearchInterface):
    """
    Extended interface that includes domain adaptability, collaboration features, and graph visualization.
    """
    def _build_main_interface(self) -> None:
        super()._build_main_interface()