import streamlit as st
import google.generativeai as genai
import json
import re

# Configurações de chave secreta e API do Google Gemini
API_KEY = st.secrets["GOOGLE_API_KEY"]  # Carrega a chave do arquivo secrets.toml
genai.configure(api_key=API_KEY)

# Configuração do modelo de IA
CONFIGURACAO_GERACAO = {
    "temperature": 0.2,
    "top_p": 0.8,
    "top_k": 40,
    "response_mime_type": "text/plain",
}
NOME_MODELO = "gemini-2.0-flash-exp"
MODELO = genai.GenerativeModel(
    model_name=NOME_MODELO,
    generation_config=CONFIGURACAO_GERACAO,
)

# Estilos CSS para o Streamlit
st.markdown(
    """
    <style>
    body {
        background-color: #f0f2f6;
        color: #262730;
    }
    .stApp {
        max-width: 900px;
        margin: 0 auto;
        padding: 20px;
    }
    .stTextInput>div>div>input {
        border-radius: 8px;
        padding: 10px;
        border: 1px solid #ced4da;
    }
    .stButton>button {
        background-color: #007bff;
        color: white;
        border: none;
        border-radius: 8px;
        padding: 10px 20px;
        cursor: pointer;
    }
    .stButton>button:hover {
        background-color: #0056b3;
    }
    .stSelectbox>div>div>div {
        border-radius: 8px;
        padding: 5px;
        border: 1px solid #ced4da;
        
    }
    .stRadio>div{
        display:flex;
        flex-direction: row;
    }
    .stRadio>div>div{
        margin-right:15px
    }
    .stDownloadButton>button{
       background-color: #28a745;
        color: white;
        border: none;
        border-radius: 8px;
        padding: 10px 20px;
        cursor: pointer;
    }
    .stDownloadButton>button:hover {
        background-color: #1e7e34;
    }

    .stAlert {
    padding: 15px;
    margin-bottom: 20px;
    border: 1px solid transparent;
    border-radius: 4px;
    }
    .stAlert.stAlert-danger {
        color: #721c24;
        background-color: #f8d7da;
        border-color: #f5c6cb;
    }
     .stSpinner {
        display: flex;
        justify-content: center;
        align-items: center;
        margin-top: 20px;
    }
    .stSpinner div {
       border: 8px solid #f3f3f3;
       border-top: 8px solid #3498db;
       border-radius: 50%;
       width: 60px;
       height: 60px;
       animation: spin 2s linear infinite;
    }

    @keyframes spin {
        0% { transform: rotate(0deg); }
        100% { transform: rotate(360deg); }
    }
    </style>
    """,
    unsafe_allow_html=True,
)

def enviar_mensagem_para_modelo(mensagem):
    """Envia uma mensagem para o modelo de IA e retorna a resposta."""
    try:
        resposta = MODELO.start_chat(history=[]).send_message(mensagem)
        return resposta.text
    except Exception as e:
        st.error(f"❌ Erro ao comunicar com a IA: {e}")
        return None

def gerar_comando_powershell(prompt_base, nivel_detalhe, tipo_script, nivel_seguranca):
    """Gera um comando Powershell com base nas configurações do usuário."""
    prompt = f"""
    Você é um especialista em Powershell. Sua tarefa é gerar um único comando Powershell com base na seguinte descrição:
    
    **Descrição do Comando:** {prompt_base}
    
    **Nível de Detalhe:** {nivel_detalhe}
    **Tipo de Script:** {tipo_script}
    **Nível de Segurança:** {nivel_seguranca}
    
    **Formato da resposta:**
    -   Responda apenas o código do comando Powershell, sem texto adicional.
    -   Não inclua comentários, explicações ou markdown.
    -   Se for gerar um comando com uma cadeia de comandos, não use ponto e virgula no final ou inicio.
    -   Se precisar use pipe "|" para encadear.
    -   Não use nenhum tipo de formatação especial no resultado, apenas o código.
    -   Se a descrição pedir para criar um arquivo, o comando deverá criar o arquivo diretamente no sistema de arquivos e não usar um output em tela para isso.
    -   Se a descrição pedir para ler um arquivo, o comando deverá ler o arquivo diretamente do sistema de arquivos e não usar um input em tela para isso.
    -   O sistema operacional padrão é Windows Server 2016 e a versão do Powershell padrão é a 7, a não ser que o usuário especifique o contrário.
    -   Certifique-se de que o comando gerado seja seguro e siga as melhores práticas do Powershell.
    
    **Importante:**
    -   Gere apenas um comando por vez.
    """
    resposta = enviar_mensagem_para_modelo(prompt)
    return resposta

def parsear_e_salvar_ps1(codigo_ia, titulo_curto):
    """Parseia o código retornado pela IA para extrair o bloco de comando ps1 e salva como .ps1."""
    # Regex para encontrar blocos de código Powershell (ignorando case)
    match = re.search(r'```powershell\s*(.*?)\s*```', codigo_ia, re.DOTALL | re.IGNORECASE)
    if match:
        codigo_ps1 = match.group(1).strip()
    else:
        codigo_ps1 = codigo_ia.strip() # Se não houver bloco, usa o código inteiro
    
    nome_arquivo = f"comando_{titulo_curto}.ps1"
    with open(nome_arquivo, "w", encoding="utf-8") as f:
        f.write(codigo_ps1)
    return nome_arquivo, codigo_ps1

def main():
    st.set_page_config(page_title="Google Gemini 2 PowerShell Command Gen", page_icon="🤖")
    st.title("🤖 Google Gemini 2 PowerShell Command Gen by [Elias Andrade](https://github.com/chaos4455)")
    st.markdown("Crie comandos PowerShell de forma fácil e rápida!")
    st.markdown("---")

    # Configurações do Usuário
    st.sidebar.header("⚙️ Configurações")
    powershell_version = st.sidebar.selectbox("Versão do PowerShell", ["7", "5.1"], index=0)
    sistema_operacional = st.sidebar.selectbox("Sistema Operacional", ["Windows Server 2016", "Windows 10", "Windows 11", "Outro"], index=0)

    nivel_detalhe = st.sidebar.selectbox("Nível de Detalhe", ["Mais detalhado", "Padrão", "Mais conciso"], index=1)
    tipo_script = st.sidebar.selectbox("Tipo de Script", ["Mais automático", "Mais interativo"], index=0)
    nivel_seguranca = st.sidebar.radio("Nível de Segurança", ["Alto", "Médio", "Baixo"], index=1)

    # Prompt base do usuário
    prompt_base = st.text_input("Descreva o Comando PowerShell:", placeholder="Ex: Liste todos os processos em execução", key="prompt_base")

    if st.button("✨ Gerar Comando PowerShell"):
        if not prompt_base:
            st.error("⚠️ Por favor, insira uma descrição do comando.")
            return

        with st.spinner("⏳ Gerando comando..."):
            codigo_ia = gerar_comando_powershell(prompt_base, nivel_detalhe, tipo_script, nivel_seguranca)
            if codigo_ia:
               
                st.markdown("### ✅ Comando Gerado:")
                st.code(codigo_ia, language="powershell")
                
                titulo_curto = prompt_base[:30].strip().replace(" ", "_").lower()
                nome_arquivo_ps1, codigo_ps1  = parsear_e_salvar_ps1(codigo_ia, titulo_curto)
                
                st.download_button(
                    label="⬇️ Baixar Comando (.ps1)",
                    data=codigo_ps1,
                    file_name=nome_arquivo_ps1,
                    mime="application/powershell",
                )
            else:
                 st.error("❌ Erro ao gerar o comando. Verifique a conexão com a IA e tente novamente.")

if __name__ == "__main__":
    main()