import gradio as gr
import os
from git import Repo
from datetime import datetime
from typing import List
from pydantic import BaseModel
import instructor
from litellm import completion
import json
import getpass
import shutil
import tempfile
import subprocess
import stat

# Initialize instructor client
client = instructor.from_litellm(completion)

class Phrases(BaseModel):
    """Response structure for the assistant"""
    phrases: List[str]

def process_questions(llm_model: str, openai_api_Key: str, deploy_key: str, metrics_prompt: str, dimensions_prompt: str, questions: str, customer_name: str, progress=gr.Progress()):
    # Check which fields are missing
    missing_fields = []
    if not llm_model: missing_fields.append("OpenAI Model")
    if not openai_api_Key: missing_fields.append("OpenAI API Key")
    if not deploy_key: missing_fields.append("GitHub Deploy Key")
    if not customer_name: missing_fields.append("Customer Name")
    if not metrics_prompt: missing_fields.append("Metrics Prompt")
    if not dimensions_prompt: missing_fields.append("Dimensions Prompt")
    if not questions: missing_fields.append("Questions")
    
    if missing_fields:
        return f"Error: The following fields are mandatory and missing: {', '.join(missing_fields)}", None
    
    # Set API key for LiteLLM
    os.environ["OPENAI_API_KEY"] = openai_api_Key
    
    # remove new lines from metrics and dimensions prompts
    metrics_prompt = metrics_prompt.replace('\n', ' ')
    dimensions_prompt = dimensions_prompt.replace('\n', ' ')
    
    intents = {
        'metrics': {
            'prompt': metrics_prompt,
        },
        'dimensions': {
            'prompt': dimensions_prompt,
        }
    }
    
    # Split questions into list
    question_list = [q.strip() for q in questions.split('\n') if q.strip()]
    
    results = {'metrics_prompt': metrics_prompt, 'dimensions_prompt': dimensions_prompt, 'results': []}
    results['llm_model'] = llm_model
    results['username'] = getpass.getuser()
    error_messages = []
    
    for idx, question in progress.tqdm(enumerate(question_list), desc="Processing questions", total=len(question_list)):
        result = {'question': question}
        try:
            for intent in intents:
                # Make API call using LiteLLM with Instructor
                response = client.chat.completions.create(
                    model=llm_model,
                    messages=[
                        {"role": "user", "content": intents[intent]['prompt'].format(question=question)}
                    ],
                    response_model=Phrases
                )
                # Store results
                result[intent] = response.phrases
                if idx == 0:
                    results[f'{intent}_prompt'] = intents[intent]['prompt']
            results['results'].append(result)
        except Exception as e:
            error_msg = f'Error processing question {question}: {str(e)}'
            print(error_msg)
            error_messages.append(error_msg)
            results['results'].append({
                "question": question,
                "error": str(e),
            })
    
    # Change working directory to current file's location
    script_dir = os.path.dirname(os.path.abspath(__file__))
    os.chdir(script_dir)
    
    # Save to file
    customer_name = customer_name.replace(' ', '_')
    output_dir = os.path.join(script_dir, "atomic_concepts_gradio")
    os.makedirs(output_dir, exist_ok=True)  # Create directory if it doesn't exist
    filename = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{customer_name}.json"
    filepath = os.path.join(output_dir, filename)
    print(f'saving output to {filepath}')
    with open(filepath, 'w') as f:
        json.dump(results, f, indent=4)
    
    try:
        # Set up SSH key before git operations
        try:
            # Save the key file and get its path
            key_path = setup_ssh_key(deploy_key)
            
            # Explicitly set the SSH command with the deploy key file path and strict host key checking disabled
            ssh_command = f'ssh -i {key_path} -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
            os.environ['GIT_SSH_COMMAND'] = ssh_command
            
            # Test SSH connection with strict host key checking disabled
            result = subprocess.run(
                ['ssh', '-T', '-i', key_path, '-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null', 'git@github.com'], 
                capture_output=True, 
                text=True
            )
            if result.returncode != 1:  # GitHub's SSH test always returns 1 when successful
                raise Exception(f"SSH test failed with return code {result.returncode}\nOutput: {result.stderr}")
        except Exception as e:
            raise Exception(f"SSH connection test failed: {str(e)}")

        # Rest of the git operations
        repo_url = "git@github.com:spotonix-inc/eval_results.git"
        temp_dir = tempfile.mkdtemp()
        
        try:
            # Clone using the SSH command with the key file
            repo = Repo.clone_from(repo_url, temp_dir, env={"GIT_SSH_COMMAND": ssh_command})
            
            # Copy the new file to the cloned repo under atomic_concepts_gradio directory
            repo_subdir = os.path.join(temp_dir, "atomic_concepts_gradio")
            os.makedirs(repo_subdir, exist_ok=True)
            shutil.copy2(filepath, os.path.join(repo_subdir, filename))
            
            # Add, commit and push from temporary directory
            repo.index.add([os.path.join("atomic_concepts_gradio", filename)])
            repo.index.commit(f"Add evaluation results: {filename}")
            origin = repo.remote('origin')
            push_info = origin.push()
            
            # Verify push was successful
            if push_info[0].flags & push_info[0].ERROR:
                raise Exception("Failed to push to repository")
            
            # Construct GitHub web URL from SSH URL
            github_web_url = repo_url.replace("git@github.com:", "https://github.com/")
            github_web_url = github_web_url.replace(".git", "")
            file_url = f"{github_web_url}/blob/main/atomic_concepts_gradio/{filename}"
            
            final_message = (
                f"Successfully processed {len(question_list)} questions.\n"
                f"Results saved to {filename} and pushed to repository.\n"
                f"View results at: {file_url}"
            )
            if error_messages:
                final_message += "\n\nErrors encountered:\n" + "\n".join(error_messages)
            return final_message, filepath
        
        except Exception as e:
            raise Exception(f"Git operation failed: {str(e)}")
        finally:
            # Clean up
            try:
                shutil.rmtree(temp_dir)
            except:
                pass
        
    except Exception as e:
        error_message = f"Error pushing to repository: {str(e)}\nResults were saved locally to: {filepath}"
        return error_message, filepath
    finally:
        # Clean up the key file
        try:
            if 'key_path' in locals():
                os.remove(key_path)
        except:
            pass

def setup_ssh_key(ssh_key_file):
    try:
        # Check if a file was uploaded
        if ssh_key_file is None:
            raise ValueError("No SSH key file was uploaded")
            
        # Save the key to the current working directory
        key_path = os.path.join(os.getcwd(), 'github_deploy_key')
        shutil.copy2(ssh_key_file.name, key_path)
        
        # Set correct permissions (600)
        os.chmod(key_path, stat.S_IRUSR | stat.S_IWUSR)
        
        return key_path
    except Exception as e:
        raise ValueError(f"Invalid SSH key: {str(e)}")

def handle_ssh_key(ssh_key):
    try:
        key_path = setup_ssh_key(ssh_key)
        # Test the SSH connection
        result = subprocess.run(['ssh', '-T', 'git@github.com'], 
                              capture_output=True, 
                              text=True)
        return f"SSH key setup complete. Key saved to {key_path}"
    except Exception as e:
        return f"Error setting up SSH key: {str(e)}"

def validate_ssh_key(key_content):
    # Check if it looks like a private key
    if not key_content.startswith('-----BEGIN'):
        raise ValueError("Invalid key format: Key must start with '-----BEGIN'")
    
    if not key_content.strip().endswith('-----END') or not key_content.strip().endswith('-----END OPENSSH PRIVATE KEY-----'):
        raise ValueError("Incomplete key: Key must end with '-----END' or '-----END OPENSSH PRIVATE KEY-----'")
    
    return True

# Create Gradio interface
with gr.Blocks(css="footer {visibility: hidden}") as iface:
    gr.Markdown("# Question Evaluation Tool")
    gr.Markdown("Enter your API keys and questions to evaluate. Questions should be separated by newlines.")
    
    with gr.Row():
        with gr.Column(scale=1):
            model = gr.Textbox(
                label="OpenAI Model",
                value="gpt-4o",
                info="Required",
                elem_id="model"
            )
            api_key = gr.Textbox(
                label="OpenAI API Key",
                type="password",
                info="Required",
                elem_id="api_key"
            )
            deploy_key = gr.File(
                label="GitHub Deploy Key",
                elem_id="deploy_key"
            )
            customer_name = gr.Textbox(
                label="Customer Name",
                info="Required",
                elem_id="customer_name"
            )
            metrics_prompt = gr.Textbox(
                label="Metrics Prompt",
                lines=10,
                max_lines=10,
                info="Required",
                elem_id="metrics_prompt"
            )
            dimensions_prompt = gr.Textbox(
                label="Dimensions Prompt",
                lines=10,
                max_lines=10,
                info="Required",
                elem_id="dimensions_prompt"
            )
        
        with gr.Column(scale=1):
            questions = gr.Textbox(label="Questions (one per line)", lines=10, max_lines=10, info="Required")
            result = gr.Textbox(label="Result", lines=10, max_lines=10)
            file_output = gr.State()  # Add this to store the filepath
            
            with gr.Row():
                process_btn = gr.Button("Process")
                download_btn = gr.Button("Download Results", interactive=False)
            
            def enable_download(message, filepath):
                # Only enable download if we have a valid filepath
                if filepath is None:
                    return message, None, gr.Button(value="Download Results", interactive=False)
                return message, filepath, gr.Button(value="Download Results", interactive=True)
            
            process_btn.click(
                fn=process_questions,
                inputs=[model, api_key, deploy_key, metrics_prompt, dimensions_prompt, questions, customer_name],
                outputs=[result, file_output],
                show_progress=True
            ).then(
                fn=enable_download,
                inputs=[result, file_output],
                outputs=[result, file_output, download_btn]
            )
            
            download_btn.click(
                fn=lambda filepath: gr.File(value=filepath, label="Download JSON") if filepath else None,
                inputs=[file_output],
                outputs=gr.File(label="Download JSON")
            )


if __name__ == "__main__":
    iface.launch(share=True)