import streamlit as st
from crewai import Agent, Task, Crew
import os
from langchain_groq import ChatGroq
from fpdf import FPDF
import pandas as pd
import plotly.express as px
import time

# Title and Sidebar
st.title("🤖 Patent Insights Consultant")

st.sidebar.write(
    "This Patent Insights Consultant uses a multi-agent system to provide actionable insights and data analysis for the patent domain."
)

# User Inputs
patent_area = st.text_input('Enter Patent Technology Area', value="Artificial Intelligence")
stakeholder = st.text_input('Enter Stakeholder', value="Patent Attorneys")

# Advanced Options
st.sidebar.subheader("Advanced Options")

# Enable/Disable Advanced Features
enable_advanced_analysis = st.sidebar.checkbox("Enable Advanced Analysis", value=True)
enable_custom_visualization = st.sidebar.checkbox("Enable Custom Visualizations", value=True)

# Optional Customization
st.sidebar.subheader("Agent Customization")

# Display customization section in a collapsible expander
with st.sidebar.expander("Customize Agent Goals", expanded=False):
    enable_customization = st.checkbox("Enable Custom Goals")
    if enable_customization:
        planner_goal = st.text_area(
            "Planner Goal",
            value="Research trends in patent filings and technological innovation, identify key players, and provide strategic recommendations."
        )
        writer_goal = st.text_area(
            "Writer Goal",
            value="Craft a professional insights document summarizing trends, strategies, and actionable outcomes for stakeholders."
        )
        analyst_goal = st.text_area(
            "Analyst Goal",
            value="Perform detailed statistical analysis of patent filings, growth trends, and innovation distribution."
        )
    else:
        planner_goal = "Research trends in patent filings and technological innovation, identify key players, and provide strategic recommendations."
        writer_goal = "Craft a professional insights document summarizing trends, strategies, and actionable outcomes for stakeholders."
        analyst_goal = "Perform detailed statistical analysis of patent filings, growth trends, and innovation distribution."

#=================
# LLM Object
#=================
llm = ChatGroq(groq_api_key=os.getenv("GROQ_API_KEY"), model="groq/llama-3.3-70b-versatile")

#=================
# Crew Agents
#=================

planner = Agent(
    role="Patent Research Consultant",
    goal=planner_goal,
    backstory=(
        "You're tasked with researching {topic} patents and identifying key trends and players. Your work supports the Patent Writer and Data Analyst."
    ),
    allow_delegation=False,
    verbose=True,
    llm=llm
)

writer = Agent(
    role="Patent Insights Writer",
    goal=writer_goal,
    backstory=(
        "Using the research from the Planner and data from the Analyst, craft a professional document summarizing patent insights for {stakeholder}."
    ),
    allow_delegation=False,
    verbose=True,
    llm=llm
)

analyst = Agent(
    role="Patent Data Analyst",
    goal=analyst_goal,
    backstory=(
        "Analyze patent filing data and innovation trends in {topic} to provide statistical insights. Your analysis will guide the Writer's final report."
    ),
    allow_delegation=False,
    verbose=True,
    llm=llm
)

#=================
# Crew Tasks
#=================

plan = Task(
    description=(
        "1. Research recent trends in {topic} patent filings and innovation.\n"
        "2. Identify key players and emerging technologies.\n"
        "3. Provide recommendations for stakeholders on strategic directions.\n"
        "4. Limit the output to 500 words."
    ),
    expected_output="A research document with structured insights and strategic recommendations.",
    agent=planner
)

write = Task(
    description=(
        "1. Use the Planner's and Analyst's outputs to craft a professional patent insights document.\n"
        "2. Include key findings, visual aids, and actionable strategies.\n"
        "3. Align content with stakeholder goals.\n"
        "4. Limit the document to 400 words."
    ),
    expected_output="A polished, stakeholder-ready patent insights document.",
    agent=writer
)

analyse = Task(
    description=(
        "1. Perform statistical analysis of patent filing trends, innovation hot spots, and growth projections.\n"
        "2. Collaborate with the Planner and Writer to align on data needs.\n"
        "3. Present findings in an actionable format."
    ),
    expected_output="A detailed statistical analysis with actionable insights for stakeholders.",
    agent=analyst
)

#=================
# Execution
#=================

crew = Crew(
    agents=[planner, analyst, writer],
    tasks=[plan, analyse, write],
    verbose=True
)

def generate_pdf_report(result):
    """Generate a professional PDF report from the Crew output."""
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    pdf.set_auto_page_break(auto=True, margin=15)

    # Title
    pdf.set_font("Arial", size=16, style="B")
    pdf.cell(200, 10, txt="Patent Insights Report", ln=True, align="C")
    pdf.ln(10)

    # Content
    pdf.set_font("Arial", size=12)
    pdf.multi_cell(0, 10, txt=result)

    # Save PDF
    report_path = "Patent_Insights_Report.pdf"
    pdf.output(report_path)
    return report_path

def create_visualizations(analyst_output):
    """Create visualizations for advanced insights."""
    if enable_custom_visualization and analyst_output:
        try:
            data = pd.DataFrame(analyst_output)
            st.markdown("### Advanced Visualization")

            fig = px.bar(data, x="Category", y="Values", title="Patent Trends by Category")
            st.plotly_chart(fig)
        except Exception as e:
            st.warning(f"Failed to create visualizations: {e}")

if st.button("Generate Patent Insights"):
    with st.spinner('Processing...'):
        try:
            start_time = time.time()
            results = crew.kickoff(inputs={"topic": patent_area, "stakeholder": stakeholder})
            elapsed_time = time.time() - start_time

            # Display Final Report (Writer's Output)
            st.markdown("### Final Report:")
            writer_output = getattr(results.tasks_output[2], "raw", "No details available.")
            if writer_output:
                st.write(writer_output)
            else:
                st.warning("No final report available.")

            # Option for Detailed Insights
            with st.expander("Explore Detailed Insights"):
                tab1, tab2 = st.tabs(["Planner's Insights", "Analyst's Analysis"])

                # Planner's Output
                with tab1:
                    st.markdown("### Planner's Insights")
                    planner_output = getattr(results.tasks_output[0], "raw", "No details available.")
                    st.write(planner_output)

                # Analyst's Output
                with tab2:
                    st.markdown("### Analyst's Analysis")
                    analyst_output = getattr(results.tasks_output[1], "raw", "No details available.")
                    st.write(analyst_output)

                    # Generate visualizations if enabled
                    if enable_advanced_analysis:
                        create_visualizations(analyst_output)

            # Display Token Usage and Execution Time
            st.success(f"Analysis completed in {elapsed_time:.2f} seconds.")
            token_usage = getattr(results, "token_usage", None)
            if token_usage:
                st.markdown("#### Token Usage")
                st.json(token_usage)

            # Generate PDF Report
            if writer_output:
                report_path = generate_pdf_report(writer_output)
                with open(report_path, "rb") as report_file:
                    st.download_button("Download Report", data=report_file, file_name="Patent_Report.pdf")

        except Exception as e:
            st.error(f"An error occurred during execution: {e}")

# Add reference and credits in the sidebar
st.sidebar.markdown("---")
st.sidebar.markdown("### Reference:")