import os
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from typing import Dict, Any
import json

def convert_code(input_code: str, model_name: str = "llama-3.3-70b-versatile") -> Dict[str, str]:
    """
    Converts input Python/IPYNB code into a Gradio app structure.
    Returns a dictionary containing app.py, requirements.txt, and README.md content.
    """
    
    # Check for API Key
    if not os.environ.get("GROQ_API_KEY"):
         raise ValueError("GROQ_API_KEY environment variable is not set.")

    chat = ChatGroq(temperature=0, model_name=model_name)

    system_prompt = """You are an expert Python developer specializing in Gradio and Hugging Face Spaces. 
    Your task is to convert the provided Python code (which might be a script or a notebook content) into a deployable Gradio web application.
    
    You must output a JSON object with exactly three keys:
    1. "app_py": The complete code for app.py. It must use Gradio to create a UI for the functionality in the source code. Ensure all imports are correct.
    2. "requirements_txt": A list of dependencies required to run the app. Include 'gradio'.
    3. "readme_md": A README.md file customized for a Hugging Face Space.
    
    Rules for app.py:
    - Encapsulate logic in functions.
    - Create a professional Gradio interface `demo = gr.Interface(...)` or `with gr.Blocks() as demo: ...`.
    - Ensure `demo.launch()` is called at the end if it's main, but standard HF spaces just look for `demo` object or run the script.
    - Handle potential errors gracefully.
    
    Do not include markdown triple backticks in the JSON values. The values should be raw string content.
    """

    human_template = "Convert this code into a Gradio app:\n\n{code}"
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        ("human", human_template)
    ])

    chain = prompt | chat | JsonOutputParser()

    try:
        result = chain.invoke({"code": input_code})
        return result
    except Exception as e:
        # Fallback or error handling
        raise RuntimeError(f"Failed to generate code: {str(e)}")

def parse_notebook(notebook_content: dict) -> str:
    """Extracts code cells from a notebook dictionary."""
    code = []
    for cell in notebook_content.get('cells', []):
        if cell.get('cell_type') == 'code':
            code.append("".join(cell.get('source', [])))
    return "\n\n".join(code)