import gradio as gr
import json
import os
from datetime import datetime
from typing import Dict, List, Optional, Union
from dataclasses import dataclass
from pathlib import Path

# Type definitions and data structures
@dataclass
class AIModel:
    name: str
    description: str
    date_of_release: str
    developer: str
    use_case: str
    impact: str
    time_difference: Optional[int] = None

    def to_dict(self) -> Dict:
        return {
            "name": self.name,
            "description": self.description,
            "dateOfRelease": self.date_of_release,
            "developer": self.developer,
            "use_case": self.use_case,
            "impact": self.impact,
            "time_difference": self.time_difference
        }

    @classmethod
    def from_dict(cls, data: Dict) -> 'AIModel':
        return cls(
            name=data.get('name', ''),
            description=data.get('description', ''),
            date_of_release=data.get('dateOfRelease', ''),
            developer=data.get('developer', ''),
            use_case=data.get('use_case', ''),
            impact=data.get('impact', ''),
            time_difference=data.get('time_difference')
        )

class ModelRepository:
    def __init__(self, file_path: str = "models.json"):
        self.file_path = Path(file_path)
        self.models: List[AIModel] = []
        self.load_models()

    def load_models(self) -> None:
        if self.file_path.exists():
            try:
                with open(self.file_path, "r") as f:
                    data = json.load(f)
                self.models = [AIModel.from_dict(model_data) for model_data in data]
            except json.JSONDecodeError:
                print(f"Error reading {self.file_path}. Starting with empty model list.")
                self.models = []
        else:
            self.models = []

    def save_models(self) -> None:
        with open(self.file_path, "w") as f:
            json.dump([model.to_dict() for model in self.models], f, indent=4)

    def add_model(self, model: AIModel) -> None:
        self.models.append(model)
        self.calculate_time_differences()
        self.save_models()

    def update_model(self, index: int, model: AIModel) -> bool:
        if 0 <= index < len(self.models):
            self.models[index] = model
            self.calculate_time_differences()
            self.save_models()
            return True
        return False

    def get_filtered_models(
        self,
        developer: Optional[str] = None,
        use_case: Optional[str] = None
    ) -> List[AIModel]:
        filtered_models = self.models.copy()

        if developer and developer != "All":
            filtered_models = [m for m in filtered_models if m.developer == developer]
        
        if use_case and use_case != "All":
            filtered_models = [m for m in filtered_models if m.use_case == use_case]

        return sorted(filtered_models, key=lambda x: x.date_of_release, reverse=True)

    def calculate_time_differences(self) -> None:
        # Sort by date in ascending order
        sorted_models = sorted(self.models, key=lambda x: x.date_of_release)
        
        # Reset all time differences
        for model in sorted_models:
            model.time_difference = None
        
        # Calculate time differences starting from the most recent
        for i in range(len(sorted_models)-1, 0, -1):
            curr_date = datetime.strptime(sorted_models[i].date_of_release, "%Y-%m-%d")
            prev_date = datetime.strptime(sorted_models[i-1].date_of_release, "%Y-%m-%d")
            sorted_models[i].time_difference = (curr_date - prev_date).days

    def get_unique_developers(self) -> List[str]:
        return sorted(list(set(model.developer for model in self.models)))

    def get_unique_use_cases(self) -> List[str]:
        return sorted(list(set(model.use_case for model in self.models)))

class UIRenderer:
    @staticmethod
    def render_model_card(model: AIModel) -> str:
        return f"""
        <div class='ai-card'>
            <h3>{model.name} ({model.date_of_release})</h3>
            <p><strong>Description:</strong> {model.description}</p>
            <p><strong>Developer:</strong> {model.developer}</p>
            <p><strong>Use Case:</strong> {model.use_case}</p>
            <p><strong>Impact:</strong> {model.impact}</p>
        </div>
        """

    @staticmethod
    def render_time_difference(days: int) -> str:
        if days <= 0:
            return ""
        gap_height = days * 2  # Scale the gap based on days
        return f"""
        <div class='time-difference' style='margin: {gap_height}px 0;'>
            <div class='time-line'></div>
            <div class='time-text'>{days} days between releases</div>
        </div>
        """

    @staticmethod
    def render_timeline(models: List[AIModel]) -> str:
        output = "<div class='timeline-container'>"
        
        for model in models:  # models are already sorted in reverse order
            if model.time_difference:  # Check if there's a time difference
                output += UIRenderer.render_time_difference(model.time_difference)
            output += UIRenderer.render_model_card(model)

        output += "</div>"
        return output

def create_ui(repository: ModelRepository) -> gr.Blocks:
    css = """
    body {
        font-family: 'Arial', sans-serif;
        background-color: #121212;
        color: white;
        margin: 0;
        padding: 0;
    }
    .timeline-container {
        display: flex;
        flex-direction: column;
        gap: 1.5rem;
        padding: 1rem;
    }
    .ai-card {
        background: linear-gradient(145deg, #2b2b2b, #333);
        border-radius: 1rem;
        padding: 1.5rem;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        border: 1px solid rgba(255, 255, 255, 0.1);
        transition: transform 0.2s, box-shadow 0.2s;
    }
    .ai-card:hover {
        transform: translateY(-2px);
        box-shadow: 0 6px 8px rgba(0, 0, 0, 0.2);
    }
    .ai-card h3 {
        color: #60a5fa;
        margin: 0 0 1rem 0;
        font-size: 1.25rem;
    }
    .ai-card p {
        color: #e5e7eb;
        margin: 0.5rem 0;
        line-height: 1.5;
    }
    .ai-card strong {
        color: #93c5fd;
    }
    .gradio-container {
        width: 80%;
        margin: auto;
        padding: 20px;
    }
    .gradio-markdown {
        color: white;
    }
    .gr-button {
        background-color: #333;
        color: white;
    }
    .gr-button:hover {
        background-color: #555;
    }
    .gr-dropdown, .gr-textbox {
        background-color: #333;
        color: white;
        border: 1px solid #555;
    }
    .time-difference {
        text-align: center;
        color: #9ca3af;
        font-style: italic;
        position: relative;
        padding: 20px 0;
    }
    .time-line {
        position: absolute;
        left: 50%;
        top: 0;
        bottom: 0;
        width: 2px;
        background: linear-gradient(to bottom, transparent, #555, transparent);
    }
    .time-text {
        background-color: #1a1a1a;
        padding: 5px 15px;
        border-radius: 15px;
        display: inline-block;
        position: relative;
        z-index: 1;
    }
    """

    with gr.Blocks(css=css) as app:
        gr.Markdown("# AI Timeline\n\nVisualize the development of AI models through an interactive timeline.")
        
        with gr.Tab("View Timeline"):
            with gr.Row():
                developer_filter = gr.Dropdown(
                    label="Filter by Developer",
                    choices=["All"] + repository.get_unique_developers(),
                    value="All"
                )
                use_case_filter = gr.Dropdown(
                    label="Filter by Use Case",
                    choices=["All"] + repository.get_unique_use_cases(),
                    value="All"
                )
            
            filter_button = gr.Button("Apply Filters", variant="primary")
            timeline_output = gr.HTML()

            def update_timeline(developer: str, use_case: str) -> str:
                filtered_models = repository.get_filtered_models(developer, use_case)
                return UIRenderer.render_timeline(filtered_models)

            filter_button.click(
                update_timeline,
                inputs=[developer_filter, use_case_filter],
                outputs=[timeline_output]
            )

        with gr.Tab("Add Model"):
            with gr.Row():
                name_input = gr.Textbox(label="Model Name", placeholder="Enter model name")
                date_input = gr.Textbox(label="Release Date (YYYY-MM-DD)", placeholder="Enter date of release")
            description_input = gr.Textbox(label="Description", placeholder="Enter a short description", lines=3)
            with gr.Row():
                developer_input = gr.Textbox(label="Developer", placeholder="Enter the developer or organization")
                use_case_input = gr.Textbox(label="Use Case", placeholder="Enter the primary use case")
            impact_input = gr.Textbox(label="Impact", placeholder="Enter the model's impact", lines=2)
            
            add_button = gr.Button("Add Model", variant="primary")
            add_output = gr.Markdown()

            def add_new_model(name, description, date, developer, use_case, impact):
                try:
                    datetime.strptime(date, "%Y-%m-%d")
                    model = AIModel(
                        name=name,
                        description=description,
                        date_of_release=date,
                        developer=developer,
                        use_case=use_case,
                        impact=impact
                    )
                    repository.add_model(model)
                    return "Model added successfully!", UIRenderer.render_timeline(repository.get_filtered_models())
                except ValueError:
                    return "Error: Invalid date format. Please use YYYY-MM-DD", ""

            add_button.click(
                add_new_model,
                inputs=[name_input, description_input, date_input,
                        developer_input, use_case_input, impact_input],
                outputs=[add_output, timeline_output]
            )

        with gr.Tab("Edit Model"):
            edit_index = gr.Number(label="Model Number", precision=0)
            with gr.Row():
                edit_name = gr.Textbox(label="New Model Name", placeholder="Enter new model name")
                edit_date = gr.Textbox(label="New Release Date", placeholder="Enter new date (YYYY-MM-DD)")
            edit_description = gr.Textbox(label="New Description", placeholder="Enter new description", lines=3)
            with gr.Row():
                edit_developer = gr.Textbox(label="New Developer", placeholder="Enter new developer")
                edit_use_case = gr.Textbox(label="New Use Case", placeholder="Enter new use case")
            edit_impact = gr.Textbox(label="New Impact", placeholder="Enter new impact", lines=2)
            
            edit_button = gr.Button("Edit Model", variant="primary")
            edit_output = gr.Markdown()

            def edit_existing_model(index, name, description, date, developer, use_case, impact):
                try:
                    datetime.strptime(date, "%Y-%m-%d")
                    model = AIModel(
                        name=name,
                        description=description,
                        date_of_release=date,
                        developer=developer,
                        use_case=use_case,
                        impact=impact
                    )
                    if repository.update_model(int(index) - 1, model):
                        return "Model updated successfully!", UIRenderer.render_timeline(repository.get_filtered_models())
                    return "Error: Invalid model number", ""
                except ValueError:
                    return "Error: Invalid date format. Please use YYYY-MM-DD", ""

            edit_button.click(
                edit_existing_model,
                inputs=[edit_index, edit_name, edit_description, edit_date,
                       edit_developer, edit_use_case, edit_impact],
                outputs=[edit_output, timeline_output]
            )

    return app

if __name__ == "__main__":
    repository = ModelRepository()
    app = create_ui(repository)
    app.launch()