import os

current_directory = os.getcwd()
os.environ['HOME'] = current_directory
print(f"New HOME directory: {os.environ['HOME']}")

import io
import requests
import zipfile
import hmac
# from metagpt.config import Config
from metagpt.software_company import generate_repo, ProjectRepo
import streamlit as st

CONFIG_FILE_PATH = 'config/config2.yaml'
if not os.path.exists(CONFIG_FILE_PATH):
    os.mkdir(os.path.dirname(CONFIG_FILE_PATH))

if not os.path.exists(CONFIG_FILE_PATH):
    with open(CONFIG_FILE_PATH, 'w') as f:
        f.write(os.environ['METAGPT_CONFIG_WITH_ANTHROPIC_API'])
# config = Config(CONFIG_FILE_PATH)

st.title("AI Software Coder")
def check_password():
    """Returns `True` if the user had the correct password."""

    def password_entered():
        """Checks whether a password entered by the user is correct."""
        if hmac.compare_digest(st.session_state["password"], os.environ['PAGE_PASSWORD']):
            st.session_state["password_correct"] = True
            del st.session_state["password"]  # Don't store the password.
        else:
            st.session_state["password_correct"] = False

    # Return True if the password is validated.
    if st.session_state.get("password_correct", False):
        return True

    # Show input for password.
    st.text_input(
        "Password", type="password", on_change=password_entered, key="password"
    )
    if "password_correct" in st.session_state:
        st.error("😕 Password incorrect")
    return False


if not check_password():
    st.stop()  # Do not continue if check_password is not True.

def generate_zip(dir_path):
    """Generates a zip file containing the input files."""
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w') as zip_file:
        for root, _, files in os.walk(dir_path):
            # Add each file to the zip archive
            for file in files:
                file_path = os.path.join(root, file)
                with open(file_path, 'rb') as f:
                    file_content = f.read()
                    zip_file.writestr(filename, file_content, arcname=os.path.relpath(file_path, dir_path))
    zip_buffer.seek(0)
    return zip_buffer

if "default" not in st.session_state:
    st.session_state["default"] = 2000
text_input = st.text_area("Enter your software development requirements here:", 
                          height=st.session_state["default"],
                          value="""The goal.

The instructions of the task for each file example, things to do with input and output.

Here are some examples.

<example>
### filename.json
Original document:
```markdown
## policy title
detailed section
- blah blah
```

Configuration:
```json
{
    "income_limit": 50
}
```
</example>

<example>
### filename.py
Original document:
```markdown
## policy title
detailed section
- blah blah
```

Configuration:
```py
class AgeCalculator:
    def __init__(self, birth_year):
        self.birth_year = birth_year

    def calculate_age(self, current_year):
        age = current_year - self.birth_year
        return self._validate_and_format_age(age)

    def _validate_and_format_age(self, age):
        if age < 0:
            raise ValueError("Invalid age calculated")
        return f"User is {age} years old"

def get_user_age(birth_year, current_year):
    calculator = AgeCalculator(birth_year)
    return calculator.calculate_age(current_year)
```
</example>

Here is the additional input from a lawer:

<input>
comments from Nick/Marc/...
</input>
""")

if st.button("Generate Code"):
    st.balloons()
    st.session_state["default"] = 200
    def download_content(url):
        response = requests.get(url)
        response.raise_for_status()
        return response.text
    
    def replace_urls_with_content(text):
        lines = text.splitlines()
        for i, line in enumerate(lines):
            if line.startswith('http://') or line.startswith('https://'):
                try:
                    content = download_content(line)
                    lines[i] = line.split('/')[-1] + '\n```\n' + content + '\n```'
                except requests.exceptions.RequestException as e:
                    print(f"Error downloading content from {line}: {e}")
        return '\n'.join(lines)

    st.info('Working on it.. A download button will show up below once work done.')
    
    text_input = replace_urls_with_content(text_input)
    repo: ProjectRepo = generate_repo(text_input)
    zip_buffer = generate_zip(repo.workdir())
    st.download_button(
        label="Download Code Repository",
        data=zip_buffer,
        file_name="output.zip",
        mime="application/zip"
    )