import threading
import time
import gradio as gr
import logging
import json
import re
import torch
import tempfile
import os
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Any, Union
from dataclasses import dataclass, field
from enum import Enum
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
from PIL import Image

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('gradio_builder.log')
    ]
)
logger = logging.getLogger(__name__)

# Constants
DEFAULT_PORT = 7860
MODEL_CACHE_DIR = Path("model_cache")
TEMPLATE_DIR = Path("templates")
TEMP_DIR = Path("temp")
DATABASE_PATH = Path("code_database.json")

# Ensure directories exist
for directory in [MODEL_CACHE_DIR, TEMPLATE_DIR, TEMP_DIR]:
    directory.mkdir(exist_ok=True, parents=True)

@dataclass
class Template:
    code: str
    description: str
    components: List[str] = field(default_factory=list)

class TemplateManager:
    def __init__(self, template_dir: Path):
        self.template_dir = template_dir
        self.templates: Dict[str, Template] = {}

    def load_templates(self):
        for file_path in self.template_dir.glob("*.json"):
            try:
                with open(file_path, 'r') as f:
                    template_data = json.load(f)
                    template = Template(**template_data)
                    self.templates[template_data['description']] = template
            except json.JSONDecodeError as e:
                logger.error(f"Error loading template from {file_path}: {e}")
            except KeyError as e:
                logger.error(f"Missing key in template file {file_path}: {e}")

    def save_template(self, name: str, template: Template) -> bool:
        file_path = self.template_dir / f"{name}.json"
        try:
            with open(file_path, 'w') as f:
                json.dump(dataclasses.asdict(template), f, indent=2)
            return True
        except Exception as e:
            logger.error(f"Error saving template to {file_path}: {e}")
            return False

    def get_template(self, name: str) -> Optional[str]:
        return self.templates.get(name, {}).get('code', "")

class RAGSystem:
    def __init__(self, model_name: str = "gpt2", device: str = "cuda" if torch.cuda.is_available() else "cpu", embedding_model="all-mpnet-base-v2"):
        self.device = device
        self.embedding_model = None
        self.code_embeddings = None
        self.index = None
        self.database = {'codes': [], 'embeddings': []}
        self.pipe = None

        try:
            self.tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=MODEL_CACHE_DIR)
            self.model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=MODEL_CACHE_DIR).to(device)
            self.pipe = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer, device=self.device)
            self.embedding_model = SentenceTransformer(embedding_model)
            self.load_database()
            logger.info("RAG system initialized successfully.")
        except Exception as e:
            logger.error(f"Error loading language model or embedding model: {e}. Falling back to placeholder generation.")

    def load_database(self):
        if DATABASE_PATH.exists():
            try:
                with open(DATABASE_PATH, 'r', encoding='utf-8') as f:
                    self.database = json.load(f)
                    self.code_embeddings = np.array(self.database['embeddings'])
                logger.info("Loaded code database from file.")
                self._build_index()
            except (json.JSONDecodeError, KeyError) as e:
                logger.error(f"Error loading code database: {e}. Creating new database.")
                self.database = {'codes': [], 'embeddings': []}
                self.code_embeddings = np.array([])
                self._build_index()
        else:
            logger.info("Code database does not exist. Creating new database.")
            self.database = {'codes': [], 'embeddings': []}
            self.code_embeddings = np.array([])
            self._build_index()

        if self.embedding_model and len(self.database['codes']) != len(self.database['embeddings']):
            logger.warning("Mismatch between number of codes and embeddings, rebuilding embeddings.")
            self.rebuild_embeddings()
        elif self.embedding_model is None:
            logger.warning ("Embeddings are not supported in this context.")

    def _build_index(self):
        if len(self.code_embeddings) > 0 and self.embedding_model:
            self.index = faiss.IndexFlatL2(self.code_embeddings.shape[1])  # L2 distance
            self.index.add(self.code_embeddings)

    def add_to_database(self, code: str):
        try:
            if self.embedding_model is None:
                raise ValueError("Embedding model not loaded.")
            embedding = self.embedding_model.encode(code)
            self.database['codes'].append(code)
            self.database['embeddings'].append(embedding.tolist())
            self.code_embeddings = np.vstack((self.code_embeddings, embedding)) if len(self.code_embeddings) > 0 else np.array([embedding])
            self.index.add(np.array([embedding]))
            self.save_database()
            logger.info(f"Added code snippet to database. Total size: {len(self.database['codes'])}.")
        except Exception as e:
            logger.error(f"Error adding to database: {e}")

    def save_database(self):
        try:
            with open(DATABASE_PATH, 'w', encoding='utf-8') as f:
                json.dump(self.database, f, indent=2)
            logger.info(f"Saved database to {DATABASE_PATH}.")
        except Exception as e:
            logger.error(f"Error saving database: {e}")

    def rebuild_embeddings(self):
        try:
            if self.embedding_model is None:
                raise ValueError("Embedding model not loaded.")
            embeddings = self.embedding_model.encode(self.database['codes'])
            self.code_embeddings = embeddings
            self.database['embeddings'] = embeddings.tolist()
            self._build_index()
            self.save_database()
            logger.info("Rebuilt and saved embeddings to the database.")
        except Exception as e:
            logger.error(f"Error rebuilding embeddings: {e}")

    def retrieve_similar_code(self, description: str, top_k: int = 3) -> List[str]:
        if self.embedding_model is None or self.index is None:
            logger.warning("Embedding model or index not available. Cannot retrieve similar code.")
            return []
        try:
            embedding = self.embedding_model.encode(description)
            D, I = self.index.search(np.array([embedding]), top_k)
            logger.info(f"Retrieved {top_k} similar code snippets for description: {description}.")
            return [self.database['codes'][i] for i in I[0]]
        except Exception as e:
            logger.error(f"Error retrieving similar code: {e}")
            return []

    def generate_code(self, description: str, template_code: str) -> str:
        retrieved_codes = self.retrieve_similar_code(description)
        prompt = f"Description: {description} Retrieved Code Snippets: {''.join([f'```python {code} ```' for code in retrieved_codes])} Template: ```python {template_code} ``` Generated Code: ```python "
        if self.pipe:
            try:
                generated_text = self.pipe(prompt, max_length=500, num_return_sequences=1)[0]['generated_text']
                generated_code = generated_text.split("Generated Code:")[1].strip().split('```')[0]
                logger.info("Code generated successfully.")
                return generated_code
            except Exception as e:
                logger.error(f"Error generating code with language model: {e}. Returning template code.")
                return template_code
        else:
            logger.warning("Text generation pipeline is not available. Returning placeholder code.")
            return f"# Placeholder code generation. Description: {description} {template_code}"

class GradioInterface:
    def __init__(self):
        self.template_manager = TemplateManager(TEMPLATE_DIR)
        self.template_manager.load_templates()
        self.rag_system = RAGSystem()

    def _extract_components(self, code: str) -> List[str]:
        components = []
        function_matches = re.findall(r'def (\w+)\(', code)  # added parenthesis for more accuracy
        components.extend(function_matches)
        class_matches = re.findall(r'class (\w+)\:', code)  # added colon for more accuracy
        components.extend(class_matches)
        logger.info(f"Extracted components: {components}")
        return components

    def _get_template_choices(self) -> List[str]:
        return list(self.template_manager.templates.keys())

    def launch(self, **kwargs):
        with gr.Blocks() as interface:
            gr.Markdown("## Code Generation Interface")
            description_input = gr.Textbox(label="Description", placeholder="Enter a description for the code you want to generate.")
            code_output = gr.Textbox(label="Generated Code", interactive=False)
            generate_button = gr.Button("Generate Code")
            template_choice = gr.Dropdown(label="Select Template", choices=self._get_template_choices(), value=None)
            save_button = gr.Button("Save as Template")
            status_output = gr.Textbox(label="Status", interactive=False)

            def generate_code_wrapper(description, template_choice):
                try:
                    template_code = self.template_manager.get_template(template_choice) if template_choice else ""
                    generated_code = self.rag_system.generate_code(description, template_code)
                    return generated_code, "Code generated successfully."
                except Exception as e:
                    return "", f"Error generating code: {e}"

            def save_template_wrapper(code, name, description):
                try:
                    if not name:
                        return code, "Template name cannot be empty."
                    if not code:
                        return code, "Code cannot be empty."

                    components = self._extract_components(code)
                    template = Template(code=code, description=name, components=components)
                    if self.template_manager.save_template(name, template):
                        self.rag_system.add_to_database(code)
                        return code, f"Template '{name}' saved successfully."
                    else:
                        return code, "Failed to save template."
                except Exception as e:
                    return code, f"Error saving template: {e}"

            generate_button.click(
                fn=generate_code_wrapper,
                inputs=[description_input, template_choice],
                outputs=[code_output, status_output]
            )

            save_button.click(
                fn=save_template_wrapper,
                inputs=[code_output, template_choice, description_input],
                outputs=[code_output, status_output]
            )

        logger.info("Launching Gradio interface.")
        interface.launch(**kwargs)

def main():
    logger.info("=== Application Startup ===")
    try:
        interface = GradioInterface()
        interface.launch(
            server_port=DEFAULT_PORT,
            share=False,
            debug=True
        )
    except Exception as e:
        logger.error(f"Application error: {e}")
        raise
    finally:
        logger.info("=== Application Shutdown ===")

if __name__ == "__main__":
    main()