import gradio as gr
from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
from core.settings import (
    MAX_LORAS, LORA_SOURCE_CHOICES, RESOLUTION_MAP, MODEL_DEFAULTS_CONFIG
)

default_vals = MODEL_DEFAULTS_CONFIG.get('Default', {})
DEFAULT_STEPS = default_vals.get('steps', 20)
DEFAULT_CFG = default_vals.get('cfg', 5.0)
DEFAULT_SAMPLER = default_vals.get('sampler_name', 'euler')
DEFAULT_SCHEDULER = default_vals.get('scheduler', 'simple')
DEFAULT_POS_PROMPT = default_vals.get('positive_prompt', '')
DEFAULT_NEG_PROMPT = default_vals.get('negative_prompt', '')


def create_base_parameter_ui(prefix, defaults=None):
    if defaults is None:
        defaults = {}
    
    components = {}
    
    with gr.Row():
        components[f'aspect_ratio_{prefix}'] = gr.Dropdown(
            label="Aspect Ratio",
            choices=list(RESOLUTION_MAP.get('sdxl', {}).keys()),
            value="1:1 (Square)",
            interactive=True,
            allow_custom_value=True
        )
    with gr.Row():
        components[f'width_{prefix}'] = gr.Number(label="Width", value=defaults.get('w', 1024), interactive=True)
        components[f'height_{prefix}'] = gr.Number(label="Height", value=defaults.get('h', 1024), interactive=True)
    with gr.Row():
        components[f'sampler_{prefix}'] = gr.Dropdown(
            label="Sampler", 
            choices=SAMPLER_CHOICES, 
            value=DEFAULT_SAMPLER if DEFAULT_SAMPLER in SAMPLER_CHOICES else (SAMPLER_CHOICES[0] if SAMPLER_CHOICES else 'euler')
        )
        components[f'scheduler_{prefix}'] = gr.Dropdown(
            label="Scheduler", 
            choices=SCHEDULER_CHOICES, 
            value=DEFAULT_SCHEDULER if DEFAULT_SCHEDULER in SCHEDULER_CHOICES else (SCHEDULER_CHOICES[0] if SCHEDULER_CHOICES else 'simple')
        )
    with gr.Row():
        components[f'steps_{prefix}'] = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=DEFAULT_STEPS)
        components[f'cfg_{prefix}'] = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.1, value=DEFAULT_CFG)
    with gr.Row():
        components[f'seed_{prefix}'] = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
        components[f'batch_size_{prefix}'] = gr.Slider(label="Batch Size", minimum=1, maximum=16, step=1, value=1)
    with gr.Row():
        components[f'guidance_{prefix}'] = gr.Slider(label="Guidance (FLUX)", minimum=1.0, maximum=10.0, step=0.1, value=3.5, visible=False, interactive=True)
        components[f'zero_gpu_{prefix}'] = gr.Number(label="ZeroGPU Duration (s)", value=None, placeholder="Default: 60s, Max: 120s", info="Optional: Set how long to reserve the GPU.")
            
    return components


def create_lora_settings_ui(prefix: str):
    components = {}
    
    lora_rows, lora_sources, lora_ids, lora_scales, lora_uploads = [], [], [], [], []

    with gr.Accordion("LoRA Settings", open=False, visible=True) as lora_accordion:
        components[f'lora_accordion_{prefix}'] = lora_accordion
        gr.Markdown("💡 **Tip:** When downloading from Civitai, please use the **Version ID**, not the Model ID. You can find the Version ID in the URL (e.g., `civitai.com/models/123?modelVersionId=456`) or under the model's download button. When downloading from Hugging Face, please use the format: `repo_id/filename.extension` or `repo_id/folder_path/filename.extension` (e.g., `Lightricks/LTX-Video/loras/ltx_lora.safetensors`).")
        components[f'lora_count_state_{prefix}'] = gr.State(1)

        for i in range(MAX_LORAS):
            with gr.Row(visible=i==0) as row:
                source = gr.Dropdown(label=f"LoRA Source {i+1}", choices=LORA_SOURCE_CHOICES, value=LORA_SOURCE_CHOICES[0], scale=1)
                lora_id = gr.Textbox(label="Civitai Version ID / HF file / Upload File", scale=2, type="text")
                scale = gr.Slider(label=f"Scale", minimum=0.0, maximum=2.0, step=0.05, value=1.0, scale=1)
                upload = gr.UploadButton(label="Upload", file_types=[".safetensors"], scale=1)
            
            lora_rows.append(row)
            lora_sources.append(source)
            lora_ids.append(lora_id)
            lora_scales.append(scale)
            lora_uploads.append(upload)
            
        with gr.Row():
            components[f'add_lora_button_{prefix}'] = gr.Button("Add LoRA", variant="secondary")
            components[f'delete_lora_button_{prefix}'] = gr.Button("Remove LoRA", variant="secondary", visible=False)

    components[f'lora_rows_{prefix}'] = lora_rows
    components[f'lora_sources_{prefix}'] = lora_sources
    components[f'lora_ids_{prefix}'] = lora_ids
    components[f'lora_scales_{prefix}'] = lora_scales
    components[f'lora_uploads_{prefix}'] = lora_uploads
    
    all_lora_components_flat = []
    for i in range(MAX_LORAS):
        all_lora_components_flat.extend([lora_sources[i], lora_ids[i], lora_scales[i], lora_uploads[i]])
    components[f'all_lora_components_flat_{prefix}'] = all_lora_components_flat
    
    return components