import random
from enum import Enum

import yaml
from docutils.nodes import reference
from nicegui import ui
from typing import TypedDict, List, Dict, Optional
from pathlib import Path
from living_agents import Character, CharacterAgent, LLMAgent
from pprint import pprint
import asyncio

from components import AsyncElement

from pathlib import Path
from typing import Optional, List
import yaml

from living_agents.datatypes import CharacterTemplate


class CharacterTemplateImportDialog(AsyncElement):
    _selected_template: Optional[CharacterTemplate]
    _templates: List[CharacterTemplate]

    async def build(self, *args, **kwargs) -> None:
        self._selected_template = None
        self._templates = []

        with self.element:
            with ui.card().classes('max-w-4xl mx-auto shadow-xl'):
                await self._header()
                await self._load_templates()
                if not self._templates:
                    await self._empty_state()
                else:
                    with ui.card_section().classes('p-2 w-full'):
                        await self._template_selection()
                        await self._template_preview()

    async def _header(self):
        """Beautiful header section"""
        with ui.card_section().classes('bg-gradient-to-r from-purple-50 to-blue-50 border-b'):
            with ui.row().classes('w-full items-center'):
                ui.icon('library_books', size='2.5rem').classes('text-purple-600')
                with ui.column().classes('flex-1 ml-4'):
                    ui.label('Import Character Template').classes('text-2xl font-bold text-gray-800')
                    ui.label('Choose a character template to create your roleplay agent').classes('text-gray-600')

    async def _load_templates(self):
        """Load templates with error handling"""
        try:
            self._templates = self.load_character_templates()
        except Exception as e:
            with ui.card_section():
                with ui.row().classes('items-center p-4 bg-red-50 border border-red-200 rounded-lg'):
                    ui.icon('error', size='2rem').classes('text-red-600')
                    with ui.column().classes('ml-3'):
                        ui.label('Failed to load templates').classes('font-bold text-red-800')
                        ui.label(str(e)).classes('text-red-600 text-sm')
            return

    async def _empty_state(self):
        """Show when no templates found"""
        with ui.card_section().classes('text-center py-12'):
            ui.icon('folder_open', size='4rem').classes('text-gray-400 mb-4')
            ui.label('No character templates found').classes('text-xl font-bold text-gray-600 mb-2')
            ui.label('Add .yml files to the character_templates directory').classes('text-gray-500')

    async def _template_selection(self):
        """Simple but rich dropdown selection"""
        ui.label('Choose Character Template').classes('text-lg font-bold text-gray-800 mb-3')

        # Prepare rich options
        options = {}
        for i, template in enumerate(self._templates):
            obs = len(template.get('observations', []))
            refl = len(template.get('reflections', []))
            plans = len(template.get('plans', []))
            total = obs + refl + plans

            # Create rich display text
            options[i] = f"📚 {template['name']} ({total} memories: {obs}obs, {refl}refl, {plans}plans)"

        # Styled select
        select = ui.select(
            options,
            label='Template',
            on_change=lambda e: self._select_template_by_index(e.value),
            clearable=False
        ).classes('w-full').props('outlined dense use-input hide-selected fill-input')

        # Custom styling for the dropdown
        select.props('menu-props="dense"')
        select.style('font-family: ui-monospace, monospace;')  # Monospace for alignment

    def _select_template_by_index(self, index: int):
        """Select template by index from dropdown"""
        if 0 <= index < len(self._templates):
            self._selected_template = self._templates[index]
            self._template_preview.refresh()

    @ui.refreshable
    async def _template_preview(self):
        """Right panel showing template details"""
        if not self._selected_template:
            await self._no_selection_state()
            return

        # Memory sections with tabs
        await self._memory_sections()
        ui.button('Create Character', icon='person_add',
                  on_click=lambda: self.submit(self._selected_template)).classes(
            'bg-gradient-to-r from-blue-500 to-purple-600 text-white px-6 py-2 w-full')

    async def _no_selection_state(self):
        """Show when no template is selected"""
        with ui.column().classes('w-full h-96 items-center justify-center'):
            ui.icon('touch_app', size='3rem').classes('text-gray-400 mb-4')
            ui.label('Select a template to preview').classes('text-xl text-gray-600 mb-2')
            ui.label('Choose from the list on the left to see the character details').classes('text-gray-500')

    async def _memory_sections(self):
        """Tabbed view of memory sections"""
        with ui.tabs().classes('w-full') as tabs:
            observations_tab = ui.tab('Observations', icon='visibility')
            reflections_tab = ui.tab('Reflections', icon='lightbulb')
            plans_tab = ui.tab('Plans', icon='track_changes')

        with ui.tab_panels(tabs, value=observations_tab).classes('w-full mt-4'):
            # Observations panel
            with ui.tab_panel(observations_tab):
                await self._render_memory_section('observations', 'blue')

            # Reflections panel
            with ui.tab_panel(reflections_tab):
                await self._render_memory_section('reflections', 'purple')

            # Plans panel
            with ui.tab_panel(plans_tab):
                await self._render_memory_section('plans', 'green')

    async def _render_memory_section(self, section_key: str, color: str):
        """Render a memory section with scrollable list"""
        memories = self._selected_template.get(section_key, [])

        if not memories:
            with ui.card().classes(f'border-{color}-200 bg-{color}-50'):
                with ui.card_section().classes('text-center py-8'):
                    ui.label(f'No {section_key} defined').classes('text-gray-600')
            return

        # Scrollable memory list
        with ui.scroll_area().classes('border border-gray-200 rounded-lg'):
            with ui.list().classes('w-full'):
                for i, memory in enumerate(memories):
                    await self._render_memory_item(memory, i, color)

    async def _render_memory_item(self, memory: str, index: int, color: str):
        """Individual memory item with nice styling"""
        with ui.item().classes('border-b border-gray-100 hover:bg-gray-50 transition-colors'):
            # with ui.item_section().classes('p-4'):
            with ui.row().classes('w-full items-start'):
                # Index badge
                ui.badge(str(index + 1), color=color).classes('mr-4 mt-1 text-xs min-w-6')

                # Memory text (with text wrapping)
                ui.label(memory).classes('flex-1 text-gray-800 leading-relaxed').style(
                    'white-space: normal; word-wrap: break-word;')

    def _select_template(self, template: CharacterTemplate):
        """Select a template and refresh the preview"""
        self._selected_template = template
        # self.main_content.refresh()

    @staticmethod
    def load_character_templates() -> List[CharacterTemplate]:
        """Load character templates from YAML files"""
        characters_dir = 'character_templates'
        characters_path = Path(characters_dir)
        character_templates: List[CharacterTemplate] = []

        if not characters_path.exists():
            raise FileNotFoundError(f"Characters directory '{characters_dir}' not found")

        if not characters_path.is_dir():
            raise ValueError(f"'{characters_dir}' is not a directory")

        # Find all YAML files
        yaml_files = list(characters_path.glob("*.yaml")) + list(characters_path.glob("*.yml"))

        if not yaml_files:
            raise ValueError(f"No YAML files found in '{characters_dir}'")

        for yaml_file in yaml_files:
            try:
                with open(yaml_file, 'r', encoding='utf-8') as file:
                    data = yaml.safe_load(file)

                required_fields = ['observations', 'reflections', 'plans']
                missing_fields = [field for field in required_fields if field not in data]

                if missing_fields:
                    print(f"Warning: File '{yaml_file.name}' missing fields: {missing_fields}")
                    continue

                # Create template
                character_templates.append(CharacterTemplate(
                    name=str(yaml_file.stem).replace('_', ' ').title(),
                    observations=data.get('observations', []),
                    reflections=data.get('reflections', []),
                    plans=data.get('plans', []),
                    yaml_file=yaml_file.name
                ))

            except Exception as e:
                print(f"Error loading '{yaml_file.name}': {e}")
                continue

        return character_templates

    def get_selected_template(self) -> Optional[CharacterTemplate]:
        """Get the currently selected template"""
        return self._selected_template


class CreationStep(Enum):
    STARTING = "starting"
    OBSERVATIONS = "observations"
    REFLECTIONS = "reflections"
    PLANS = "plans"
    OVERVIEW = "overview"
    COMPLETE = "complete"


class CharacterCreationDialog(AsyncElement):
    _template: CharacterTemplate
    _character: Optional[Character]
    _character_agent: CharacterAgent
    _current_step: CreationStep
    _steps_completed: set[CreationStep]
    _creation_task: Optional[asyncio.Task]

    async def build(self, template: CharacterTemplate, *args, **kwargs) -> None:
        self._template = template
        print(template)
        self._character = None
        self._current_step = CreationStep.STARTING
        self._steps_completed = set()
        self._creation_task = None

        with self.element:
            with ui.card().classes('w-96 mx-auto shadow-lg'):
                await self.character_creation_view()

        # Start creation automatically
        self._creation_task = asyncio.create_task(self._create_character_internal())

    @ui.refreshable
    async def character_creation_view(self):
        if self._character:
            await self._completion_view()
            return

        await self._progress_view()

    async def _progress_view(self):
        # Header
        with ui.row().classes('w-full items-center mb-6 p-4 bg-gradient-to-r from-blue-50 to-purple-50 rounded-lg'):
            ui.icon('psychology', size='2rem').classes('text-blue-600')
            with ui.column().classes('flex-1 ml-3'):
                ui.label(f'Creating {self._template["name"]}').classes('text-xl font-bold text-gray-800')
                ui.label('Building character memories and personality...').classes('text-sm text-gray-600')

        # Progress bar
        steps = [CreationStep.OBSERVATIONS, CreationStep.REFLECTIONS, CreationStep.PLANS, CreationStep.OVERVIEW]
        progress_value = (len(self._steps_completed) / len(steps)) * 100
        ui.linear_progress(value=progress_value / 100).classes('mb-6')
        ui.label(f'{len(self._steps_completed)}/{len(steps)} steps complete').classes(
            'text-center text-sm text-gray-600 mb-4')

        # Steps list
        steps_info = [
            (CreationStep.OBSERVATIONS, 'visibility', 'Loading Observations',
             'Injecting basic memories and experiences'),
            (CreationStep.REFLECTIONS, 'lightbulb', 'Processing Reflections',
             'Generating deeper insights and understanding'),
            (CreationStep.PLANS, 'track_changes', 'Installing Plans', 'Setting goals and future intentions'),
            (CreationStep.OVERVIEW, 'summarize', 'Extracting Profile', 'Creating character overview with AI'),
        ]

        with ui.list().classes('w-full space-y-2'):
            for step, icon, title, description in steps_info:
                await self._render_step_item(step, icon, title, description)

    async def _render_step_item(self, step: CreationStep, icon: str, title: str, description: str):
        # Status logic
        if step in self._steps_completed:
            status = 'complete'
            icon_color = 'text-green-600'
            bg_color = 'bg-green-50 border-green-200'
        elif step == self._current_step:
            status = 'current'
            icon_color = 'text-blue-600'
            bg_color = 'bg-blue-50 border-blue-200'
        else:
            status = 'pending'
            icon_color = 'text-gray-400'
            bg_color = 'bg-gray-50 border-gray-200'

        with ui.item().classes(f'rounded-lg border-2 {bg_color} transition-all duration-300'):
            with ui.row().classes('w-full items-center p-3'):

                # Icon/Spinner
                if status == 'complete':
                    ui.icon('check_circle').classes(f'{icon_color} text-2xl')
                elif status == 'current':
                    ui.spinner(size='lg').classes('text-blue-600')
                else:
                    ui.icon(icon).classes(f'{icon_color} text-2xl')

                # Text
                with ui.column().classes('flex-1 ml-4'):
                    ui.label(title).classes(
                        f'font-semibold {"text-green-800" if status == "complete" else "text-blue-800" if status == "current" else "text-gray-600"}')
                    ui.label(description).classes('text-sm text-gray-600 mt-1')

                # Badge
                if status == 'complete':
                    ui.badge('Done', color='positive').classes('ml-auto')
                elif status == 'current':
                    ui.badge('Processing...', color='info').classes('ml-auto animate-pulse')

    async def _completion_view(self):
        with ui.column().classes('w-full items-center p-6'):
            # Success icon
            with ui.row().classes('items-center justify-center mb-6 p-4 bg-green-50 rounded-full w-20 h-20 mx-auto'):
                ui.icon('celebration', size='3rem').classes('text-green-600')

            ui.label(f'{self._character.name} is ready!').classes('text-2xl font-bold text-green-800 text-center mb-2')
            ui.label('Character creation completed successfully').classes('text-gray-600 text-center mb-6')

            # Character preview
            with ui.card().classes('w-full bg-gradient-to-br from-green-50 to-blue-50 border border-green-200'):
                with ui.card_section():
                    with ui.row().classes('items-start gap-4'):
                        with ui.avatar().classes('bg-gradient-to-br from-blue-400 to-purple-600 text-white text-xl'):
                            ui.label(self._character.name[0].upper())

                        with ui.column().classes('flex-1'):
                            ui.label(f'{self._character.name}, {self._character.age}').classes('font-bold text-lg')
                            ui.label(self._character.occupation).classes('text-blue-600 font-medium')
                            ui.label(self._character.personality).classes('text-gray-700 text-sm mt-1')

            # Done button
            with ui.row().classes('gap-3 mt-6'):
                ui.button('Start Chatting', icon='chat').classes(
                    'bg-gradient-to-r from-blue-500 to-purple-600 text-white')

    async def _create_character_internal(self):
        """Mock character creation with realistic timing"""
        try:
            # self._character = Character(name=self._template['name'])
            # self._character_agent = CharacterAgent()
            # Step 1: Observations
            self._current_step = CreationStep.OBSERVATIONS
            self.character_creation_view.refresh()
            await asyncio.sleep(2.0)  # Simulate loading observations

            self._steps_completed.add(CreationStep.OBSERVATIONS)
            self.character_creation_view.refresh()
            await asyncio.sleep(0.3)

            # Step 2: Reflections
            self._current_step = CreationStep.REFLECTIONS
            self.character_creation_view.refresh()
            await asyncio.sleep(2.8)  # Reflections take longer

            self._steps_completed.add(CreationStep.REFLECTIONS)
            self.character_creation_view.refresh()
            await asyncio.sleep(0.3)

            # Step 3: Plans
            self._current_step = CreationStep.PLANS
            self.character_creation_view.refresh()
            await asyncio.sleep(1.5)

            self._steps_completed.add(CreationStep.PLANS)
            self.character_creation_view.refresh()
            await asyncio.sleep(0.3)

            # Step 4: AI Overview Extraction
            self._current_step = CreationStep.OVERVIEW
            self.character_creation_view.refresh()
            await asyncio.sleep(3.2)  # AI takes longest

            # Mock character creation
            self._character = self._create_mock_character()

            self._steps_completed.add(CreationStep.OVERVIEW)
            self._current_step = CreationStep.COMPLETE

            self.character_creation_view.refresh()

        except Exception as e:
            ui.notify(f'Error creating character: {str(e)}', color='negative')

    def _create_mock_character(self) -> Character:
        """Create mock character based on template"""
        mock_personalities = [
            "Shy and thoughtful with hidden depths",
            "Confident and outgoing tech enthusiast",
            "Creative and energetic with artistic flair"
        ]

        return Character(
            name=self._template["name"],
            age=random.randint(22, 35),
            personality=random.choice(mock_personalities),
            occupation=f"{self._template['name']}'s profession",
            location="Coffee shop"
        )

    def get_created_character(self) -> Optional[Character]:
        return self._character