import os
import gradio as gr
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# List of programming languages
LANGUAGES = [
    'Python', 'JavaScript', 'Java', 'C++', 'Ruby', 
    'Go', 'Rust', 'Swift', 'Kotlin', 'TypeScript', 
    'PHP', 'C#', 'Scala', 'Dart', 'Perl'
]

class CodeConverter:
    def __init__(self):
        """
        Initialize the Groq language model and prompt template
        """
        # Initialize Groq ChatModel
        self.llm = ChatGroq(
            temperature=0,
            model_name="llama3-groq-8b-8192-tool-use-preview",
            groq_api_key=os.getenv('GROQ_API_KEY')
        )

        # Create a prompt template for code conversion
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", """You are an expert code translator. 
            Your task is to convert code from one programming language to another.
            
            Guidelines:
            - Preserve the original logic and structure
            - Maintain the same functionality
            - Use idiomatic code for the target language
            - Handle any language-specific nuances
            - Provide clean, readable, and efficient code
            
            If direct translation is not possible, provide the most equivalent implementation."""),
            ("human", """Convert the code from {source_language} to {target_language}:

Code to convert:
```{source_language}
{input_code}
```

Please provide the converted code in {target_language}.""")
        ])

    def convert_code(self, input_code, source_lang, target_lang):
        """
        Convert code between programming languages
        """
        try:
            # Limit input code to 200 lines
            input_lines = input_code.split('\n')
            if len(input_lines) > 200:
                return "Error: Code exceeds 200 lines limit"
            
            # Create the chain
            chain = self.prompt | self.llm
            
            # Generate the conversion
            result = chain.invoke({
                "source_language": source_lang,
                "target_language": target_lang,
                "input_code": input_code
            })
            
            # Extract and clean the converted code
            converted_code = result.content.strip()
            
            # Remove potential code block markdown
            if converted_code.startswith('```') and converted_code.endswith('```'):
                converted_code = converted_code.strip('```').strip()
            
            return converted_code
        
        except Exception as e:
            return f"An error occurred: {str(e)}"

def create_code_converter():
    """
    Create Gradio interface for code conversion
    """
    # Initialize code converter
    converter = CodeConverter()
    
    # Create Gradio interface
    demo = gr.Interface(
        fn=converter.convert_code,
        inputs=[
            gr.Textbox(
                label="Enter your code", 
                lines=10, 
                placeholder="Paste your code here (max 200 lines)"
            ),
            gr.Dropdown(
                choices=LANGUAGES, 
                label="Source Language"
            ),
            gr.Dropdown(
                choices=LANGUAGES, 
                label="Target Language"
            )
        ],
        outputs=gr.Textbox(
            label="Converted Code", 
            lines=10
        ),
        title="🖥️ Code Converter",
        description="Convert code between different programming languages using Groq and Langchain"
    )
    
    return demo

# Launch the Gradio app
if __name__ == "__main__":
    converter = create_code_converter()
    converter.launch(share=True)