import streamlit as st
import os
import re
import base64
import ollama

def query_llama2(query):
    """
    Queries the Llama model and returns the response in chunks.
    
    Args:
    query (str): The query to send to the Llama model.

    Yields:
    str: The response chunk from the Llama model.
    """
    stream = ollama.chat(
        model='llama3:8b',
        messages=[{'role': 'user', 'content': query}],
        stream=True,
    )

    # Yield each chunk as it arrives
    for chunk in stream:
        chunk_content = chunk['message']['content']
        yield chunk_content

def read_file(file_path):
    """
    Reads the content of a text file and returns it as a string.

    Args:
    file_path (str): The path to the text file.

    Returns:
    str: The content of the file as a string.
    """
    try:
        with open(file_path, 'r') as file:
            content = file.read()
        return content
    except FileNotFoundError:
        return f"File '{file_path}' not found."
    except Exception as e:
        return f"An error occurred: {e}"

def get_response(report_file):
    """
    Reads a planet report from a file and queries the Llama model based on it.

    Args:
    report_file (str): The path to the planet report text file.

    Returns:
    generator: A generator that yields chunks of the response from the Llama model.
    """
    report_data = read_file(report_file)

    user_query = f'''
    {report_data}
    Based on the above information about my planet report, tell me about my personality and other traits that can be predicted from your knowledge.
    '''

    # Query the Llama model and return the response chunks
    response_generator = query_llama2(user_query)
    return response_generator, report_data

def read_svg(path_svg):
    """Read SVG file and return its content."""
    try:
        with open(path_svg, "r") as file:
            svg_logo = file.read().splitlines()
            _maped_list = map(str, svg_logo)
            svg_logo = "".join(_maped_list)
            temp_svg_logo = re.findall("<svg.*</svg>", svg_logo, flags=re.IGNORECASE)
            svg_logo = temp_svg_logo[0]
    except Exception as e:
        svg_logo = '<svg xmlns="http://www.w3.org/2000/svg" width="150px" height="1px" viewBox="0 0 150 1"></svg>'
    return svg_logo

def render_svg(svg):
    """Render SVG content in Streamlit."""
    b64 = base64.b64encode(svg.encode("utf-8")).decode("utf-8")
    html = (
        r"""
<div align="center">
<img src="data:image/svg+xml;base64,%s" alt="SVG Image" style="width: 50em;"/>
</div>
        """
        % b64
    )
    st.markdown(html, unsafe_allow_html=True)


# Streamlit application title
st.title("Astrological Report Generator")

# Create a form for user input
with st.form(key='astrology_form'):
    name = st.text_input("Name")
    year = st.number_input("Year Of Birth", min_value=1900, max_value=2100)
    month = st.number_input("Birth Month", min_value=1, max_value=12)
    day = st.number_input("Birth Date", min_value=1, max_value=31)
    hour = st.number_input("Time of Birth in hours (Don't prefix 0)", min_value=0, max_value=23)
    minute = st.number_input("Time of birth in minutes", min_value=0, max_value=59)
    city = st.text_input("Birth City")
    nation = st.text_input("Birth Nation")
    
    submit_button = st.form_submit_button("Generate Report")

if submit_button:
    # Validate input
    if not all([name, year, month, day, hour, minute, city, nation]):
        st.error("All fields are required")
    else:
        # Get the first part of the name up to the first space
        file_name = name.split()[0]
        svg_file_name = f"{name} - Natal Chart.svg"

        # Construct the command to run the other Python script
        command = f'python3 planet.py "{name}" {year} {month} {day} {hour} {minute} "{city}" "{nation}" > {file_name}.txt'

        # Execute the command using os.system()
        os.system(command)

        # Get response from chatbot
        response_generator, report = get_response(f"{file_name}.txt")
        
        # Clean up the generated file
        os.remove(f"{file_name}.txt")

        # Display the report data
        st.subheader("Birth Chart:")
        # Display the SVG file
        if os.path.exists(svg_file_name):
            svg_content = read_svg(svg_file_name)
            render_svg(svg_content)
        else:
            st.error("SVG file not found.")

        # Display the response in chunks, but only print the full response at the end
        st.subheader("Birth Chart Analysis:")
        response_container = st.empty()  # Create an empty container for dynamic updates
        
        response_chunks = []
        for chunk in response_generator:
            response_chunks.append(chunk)
            response_container.text("".join(response_chunks))  # Update the container with the latest chunks

        # After all chunks have been printed, clear the container and display the final response
        full_response = "".join(response_chunks)
        response_container.empty()  # Clear the container
        st.write(full_response)  # Display the full response
        
        # Remove the SVG file
        os.remove(svg_file_name)