import streamlit as st
from agno.agent import Agent
from agno.knowledge.pdf import PDFKnowledgeBase, PDFReader
from agno.vectordb.qdrant import Qdrant
from agno.tools.duckduckgo import DuckDuckGoTools
from agno.models.openai import OpenAIChat
from agno.embedder.openai import OpenAIEmbedder
import tempfile
import os
from agno.document.chunking.document import DocumentChunking


def init_session_state():
    """Initialize session state variables"""
    if 'openai_api_key' not in st.session_state:
        st.session_state.openai_api_key = None
    if 'qdrant_api_key' not in st.session_state:
        st.session_state.qdrant_api_key = None
    if 'qdrant_url' not in st.session_state:
        st.session_state.qdrant_url = None
    if 'vector_db' not in st.session_state:
        st.session_state.vector_db = None
    if 'legal_team' not in st.session_state:
        st.session_state.legal_team = None
    if 'knowledge_base' not in st.session_state:
        st.session_state.knowledge_base = None
    if 'processed_files' not in st.session_state:
        st.session_state.processed_files = set()

COLLECTION_NAME = "legal_documents"

def init_qdrant():
    """Initialize Qdrant client with configured settings."""
    if not all([st.session_state.qdrant_api_key, st.session_state.qdrant_url]):
        return None
    
    try:
        vector_db = Qdrant(
            collection=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhY2Nlc3MiOiJtIn0.lvNzlvZfmBoaR2g9d3La7ybdts4OdZcPsXB-R2bQ498,
            url=st.session_state.qdrant_url,
            api_key=st.session_state.qdrant_api_key,
            embedder=OpenAIEmbedder(
                id="text-embedding-3-small",
                api_key=st.session_state.openai_api_key
            )
        )
        return vector_db
    except Exception as e:
        st.error(f"🔴 Qdrant connection failed: {str(e)}")
        return None
        
def process_document(uploaded_file, vector_db: Qdrant):
    """Process document, create embeddings and store in Qdrant vector database"""
    if not st.session_state.openai_api_key:
        raise ValueError("OpenAI API key not provided")
    
    os.environ['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'] = st.session_state.openai_api_key
    
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as temp_file:
            temp_file.write(uploaded_file.getvalue())
            temp_file_path = temp_file.name
        
        st.info("Loading and processing document...")
        
        knowledge_base = PDFKnowledgeBase(
            path=temp_file_path,
            vector_db=vector_db,
            reader=PDFReader(),
            chunking_strategy=DocumentChunking(
                chunk_size=1000,
                overlap=200
            )
        )
        
        with st.spinner('📤 Loading documents into knowledge base...'):
            knowledge_base.load(recreate=True, upsert=True)
            st.success("✅ Documents stored successfully!")
        
        os.unlink(temp_file_path)
        return knowledge_base
    
    except Exception as e:
        st.error(f"Document processing error: {str(e)}")
        raise Exception(f"Error processing document: {str(e)}")
        
# Initialize agents within the document processing
legal_researcher = Agent(
    name="Legal Researcher",
    role="Legal research specialist",
    model=OpenAIChat(id="gpt-4"),
    tools=[DuckDuckGoTools()],
    knowledge=st.session_state.knowledge_base,
    search_knowledge=True,
    instructions=[
        "Find and cite relevant legal cases and precedents",
        "Provide detailed research summaries with sources",
        "Reference specific sections from the uploaded document",
        "Always search the knowledge base for relevant information"
    ],
    show_tool_calls=True,
    markdown=True
)

contract_analyst = Agent(
    name="Contract Analyst",
    role="Contract analysis specialist", 
    model=OpenAIChat(id="gpt-4"),
    knowledge=st.session_state.knowledge_base,
    search_knowledge=True,
    instructions=[
        "Review contracts thoroughly",
        "Identify key terms and potential issues",
        "Reference specific clauses from the document"
    ],
    markdown=True
)

legal_strategist = Agent(
    name="Legal Strategist",
    role="Legal strategy specialist",
    model=OpenAIChat(id="gpt-4"),
    knowledge=st.session_state.knowledge_base,
    search_knowledge=True,
    instructions=[
        "Develop comprehensive legal strategies",
        "Provide actionable recommendations", 
        "Consider both risks and opportunities"
    ],
    markdown=True
)
Create the agent team coordinator:

st.session_state.legal_team = Agent(
    name="Legal Team Lead",
    role="Legal team coordinator",
    model=OpenAIChat(id="gpt-4"),
    team=[legal_researcher, contract_analyst, legal_strategist],
    knowledge=st.session_state.knowledge_base,
    search_knowledge=True,
    instructions=[
        "Coordinate analysis between team members",
        "Provide comprehensive responses",
        "Ensure all recommendations are properly sourced",
        "Reference specific parts of the uploaded document",
        "Always search the knowledge base before delegating tasks"
    ],
    show_tool_calls=True,
    markdown=True
)
Build the Streamlit interface:

def main():
    st.set_page_config(page_title="Legal Document Analyzer", layout="wide")
    init_session_state()
    
    st.title("AI Legal Agent Team 👨‍⚖️")
    
    with st.sidebar:
        st.header("🔑 API Configuration")
        
        openai_key = st.text_input("OpenAI API Key", type="password")
        if openai_key:
            st.session_state.openai_api_key = openai_key
            
        qdrant_key = st.text_input("Qdrant API Key", type="password")
        if qdrant_key:
            st.session_state.qdrant_api_key = qdrant_key
            
        qdrant_url = st.text_input("Qdrant URL")
        if qdrant_url:
            st.session_state.qdrant_url = qdrant_url
            
analysis_configs = {
    "Contract Review": {
        "query": "Review this contract and identify key terms, obligations, and potential issues.",
        "agents": ["Contract Analyst"],
        "description": "Detailed contract analysis focusing on terms and obligations"
    },
    "Legal Research": {
        "query": "Research relevant cases and precedents related to this document.",
        "agents": ["Legal Researcher"], 
        "description": "Research on relevant legal cases and precedents"
    },
    "Risk Assessment": {
        "query": "Analyze potential legal risks and liabilities in this document.",
        "agents": ["Contract Analyst", "Legal Strategist"],
        "description": "Combined risk analysis and strategic assessment"
    },
    "Compliance Check": {
        "query": "Check this document for regulatory compliance issues.",
        "agents": ["Legal Researcher", "Contract Analyst", "Legal Strategist"],
        "description": "Comprehensive compliance analysis"
    },
    "Custom Query": {
        "query": None,
        "agents": ["Legal Researcher", "Contract Analyst", "Legal Strategist"],
        "description": "Custom analysis using all available agents"
    }
}