import streamlit as st
from crewai import Agent, Task, Crew, Process
from langchain_openai import ChatOpenAI
from langchain_groq import ChatGroq
from langchain_anthropic import ChatAnthropic
from dotenv import load_dotenv
import os
from crewai_tools import *


load_dotenv()

def create_lmstudio_llm(model, temperature):
    api_base = os.getenv('LMSTUDIO_API_BASE')
    os.environ["OPENAI_API_KEY"] = "lm-studio"
    os.environ["OPENAI_API_BASE"] = api_base
    if api_base:
        return ChatOpenAI(openai_api_key='lm-studio', openai_api_base=api_base, temperature=temperature)
    else:
        raise ValueError("LM Studio API base not set in .env file")

def create_openai_llm(model, temperature):
    safe_pop_env_var('OPENAI_API_KEY')
    safe_pop_env_var('OPENAI_API_BASE')
    load_dotenv(override=True)
    api_key = os.getenv('OPENAI_API_KEY')
    api_base = os.getenv('OPENAI_API_BASE', 'https://api.openai.com/v1/')
    if api_key:
        return ChatOpenAI(openai_api_key=api_key, openai_api_base=api_base, model_name=model, temperature=temperature)
    else:
        raise ValueError("OpenAI API key not set in .env file")

def create_groq_llm(model, temperature):
    api_key = os.getenv('GROQ_API_KEY')
    if api_key:
        return ChatGroq(groq_api_key=api_key, model_name=model, temperature=temperature)
    else:
        raise ValueError("Groq API key not set in .env file")

def create_anthropic_llm(model, temperature):
    api_key = os.getenv('ANTHROPIC_API_KEY')
    if api_key:
        return ChatAnthropic(anthropic_api_key=api_key, model_name=model, temperature=temperature)
    else:
        raise ValueError("Anthropic API key not set in .env file")

def safe_pop_env_var(key):
    try:
        os.environ.pop(key)
    except KeyError:
        pass
        
LLM_CONFIG = {
    "OpenAI": {
        "create_llm": create_openai_llm
    },
    "Groq": {
        "create_llm": create_groq_llm
    },
    "LM Studio": {
        "create_llm": create_lmstudio_llm
    },
    "Anthropic": {
        "create_llm": create_anthropic_llm
    }
}

def create_llm(provider_and_model, temperature=0.1):
    provider, model = provider_and_model.split(": ")
    create_llm_func = LLM_CONFIG.get(provider, {}).get("create_llm")
    if create_llm_func:
        return create_llm_func(model, temperature)
    else:
        raise ValueError(f"LLM provider {provider} is not recognized or not supported")

def load_agents():
    agents = [
        
Agent(
    role="Rechtschreibw\u00e4chter",
    backstory="Du bist Experte in der Deutschen Rechtschreibung.",
    goal="Finde Rechtschreibfehler und begr\u00fcnde die Fehler anhand dem Regelwerk.",
    allow_delegation=False,
    verbose=True,
    tools=[PDFSearchTool(pdf="/opt/CrewAI-Studio/pdf/RfdR_Amtliches-Regelwerk_2024.pdf")],
    llm=create_llm("OpenAI: gpt-4o-mini", 0.5)
)
            
    ]
    return agents

def load_tasks(agents):
    tasks = [
        
Task(
    description="Finde Rechtschreibfehler im Text und begr\u00fcnde, weshalb ein Fehler vorliegt anhand des Regelwerkes.",
    expected_output="Eine tabellarische \u00dcbersicht mit einer Spalte mit  dem Fehler und der zweiten Spalte mit dem Hinweis, gegen welche Regel versto\u00dfen wurde.",
    agent=next(agent for agent in agents if agent.role == "Rechtschreibw\u00e4chter"),
    async_execution=True
)
            
    ]
    return tasks

def main():
    st.title("Schulentwicklung")

    agents = load_agents()
    tasks = load_tasks(agents)
    crew = Crew(
        agents=agents, 
        tasks=tasks, 
        process="sequential", 
        verbose=True, 
        memory=True, 
        cache=True, 
        max_rpm=1000,
        
    )

    

    placeholders = {
        
    }

    if st.button("Run Crew"):
        with st.spinner("Running crew..."):
            try:
                result = crew.kickoff(inputs=placeholders)
                if isinstance(result, dict):
                    with st.expander("Final output", expanded=True):                
                        st.write(result.get('final_output', 'No final output available'))
                    with st.expander("Full output", expanded=False):
                        st.write(result)
                else:
                    st.write("Result:")
                    st.write(result)
            except Exception as e:
                st.error(f"An error occurred: {str(e)}")

if __name__ == '__main__':
    main()