"""ProjectCraft — AI-powered project management tool. Entry point."""
import asyncio, json, os
import gradio as gr
from agents import orchestrator
from views import gantt, raci, risks, heatmap
from core.planner import ProjectPlan
from export.xlsx_export import export_xlsx

# ---------------------------------------------------------------------------
# Color palette
# ---------------------------------------------------------------------------
COLORS = {
    "bg":           "#0a0a0a",
    "surface":      "#111111",
    "surface_alt":  "#181818",
    "green":        "#4ade80",
    "green_accent": "#22c55e",
    "silver":       "#e8e8e8",
    "silver_muted": "#a0a0a0",
    "border":       "#2a2a2a",
}

# ---------------------------------------------------------------------------
# Template loader
# ---------------------------------------------------------------------------
import functools


@functools.lru_cache(maxsize=None)
def _load_templates() -> dict:
    """Load all template JSON files from templates/ directory."""
    templates = {}
    templates_dir = os.path.join(os.path.dirname(__file__), "templates")
    template_files = [
        "software_launch.json",
        "marketing_campaign.json",
        "hr_onboarding.json",
        "event_planning.json",
    ]
    for filename in template_files:
        filepath = os.path.join(templates_dir, filename)
        try:
            with open(filepath, "r", encoding="utf-8") as f:
                data = json.load(f)
                name = data.get(
                    "name",
                    filename.replace(".json", "").replace("_", " ").title(),
                )
                templates[name] = data
        except FileNotFoundError:
            pass
        except (json.JSONDecodeError, Exception):
            pass
    return templates


# ---------------------------------------------------------------------------
# Initial state
# ---------------------------------------------------------------------------
INITIAL_STATE = {
    "inputs": {
        "objective": "", "start_date": "", "end_date": "",
        "start_date_ai": False, "end_date_ai": False,
        "tasks": "", "teams": [], "template": None,
        "effective_start": "", "effective_end": ""
    },
    "plan": {
        "phases": [], "raci": [], "risks": [], "workload": [],
        "estimated_start_date": None, "estimated_end_date": None,
        "timeline_note": None, "timeline_warning": None
    },
    "original_plan": None,
    "chat_history": []
}

# ---------------------------------------------------------------------------
# CSS (adaptive theme: dark by default, light when OS prefers it)
# ---------------------------------------------------------------------------
CSS = """
/* === Design tokens: dark defaults === */
:root {
  --pc-bg:           #0a0a0a;
  --pc-surface:      #111111;
  --pc-surface-alt:  #181818;
  --pc-text:         #e8e8e8;
  --pc-text-muted:   #a0a0a0;
  --pc-border:       #4ade80;
  --pc-green:        #4ade80;
  --pc-green-hover:  #22c55e;
  color-scheme: dark;
}

/* === Light mode overrides === */
@media (prefers-color-scheme: light) {
  :root {
    --pc-bg:           #f5f5f5;
    --pc-surface:      #ffffff;
    --pc-surface-alt:  #ebebeb;
    --pc-text:         #111111;
    --pc-text-muted:   #555555;
    --pc-border:       #16a34a;
    --pc-green:        #16a34a;
    --pc-green-hover:  #15803d;
    color-scheme: light;
  }
}

/* === Base === */
body, .gradio-container {
  background: var(--pc-bg) !important;
  color: var(--pc-text) !important;
}

/* === All text / labels (Gradio uses many class combos) === */
label,
.label-wrap,
.label-wrap span,
.block > label,
.block > label > span,
.svelte-1ipelgc,
span.svelte-1ipelgc,
.gradio-container p,
.gradio-container li,
.gradio-container span,
.gradio-container h1,
.gradio-container h2,
.gradio-container h3,
.prose, .prose p,
.markdown, .markdown p,
.chatbot .message,
.chatbot .message p {
  color: var(--pc-text) !important;
}
.prose h1, .markdown h1 {
  color: var(--pc-green) !important;
  font-size: 1.8rem !important;
  font-weight: 800 !important;
  letter-spacing: -0.5px;
}
.prose h2, .markdown h2 {
  color: var(--pc-green) !important;
  font-size: 1.25rem !important;
  font-weight: 700 !important;
}
.prose h3, .markdown h3 {
  color: var(--pc-green) !important;
  font-size: 1.05rem !important;
  font-weight: 700 !important;
}

/* === Inputs === */
textarea,
input[type="text"], input[type="number"],
input[type="email"], input[type="search"] {
  background: var(--pc-surface) !important;
  color: var(--pc-text) !important;
  border-color: var(--pc-border) !important;
}
textarea::placeholder,
input::placeholder { color: var(--pc-text-muted) !important; }

select {
  background: var(--pc-surface) !important;
  color: var(--pc-text) !important;
  border-color: var(--pc-border) !important;
}
input[type="checkbox"] { accent-color: var(--pc-green); }

/* === Buttons === */
button.primary, .primary {
  background: var(--pc-green) !important;
  color: #ffffff !important;
  font-weight: 700;
}
button.primary:hover, .primary:hover {
  background: var(--pc-green-hover) !important;
}
button:not(.primary) {
  background: var(--pc-surface-alt) !important;
  color: var(--pc-text) !important;
  border-color: var(--pc-border) !important;
}

/* === Tabs === */
.tab-nav button, button[role="tab"] {
  color: var(--pc-text) !important;
  font-size: 15px !important;
  font-weight: 600 !important;
}
.tab-nav button.selected, button[role="tab"][aria-selected="true"] {
  color: var(--pc-green) !important;
  border-color: var(--pc-green) !important;
  background: var(--pc-surface-alt) !important;
}

/* === Plotly SVG text — axis labels and zoom buttons bright in dark mode === */
/* Horizontal bar charts (Gantt) render categorical y-axis labels as direct
   <text> children of yaxislayer-above — NOT inside ytick groups. So [class*="tick"]
   works for numeric/date x-axis but misses categorical y-axis labels.
   [class*="axislayer"] and [class*="yaxis"] cover those containers.
   :root prefix raises specificity to 0,3,1 to beat any conflicting Gradio rules. */
:root .js-plotly-plot [class*="tick"] text,
:root .js-plotly-plot [class*="tick"] tspan,
:root .js-plotly-plot [class*="axislayer"] text,
:root .js-plotly-plot [class*="axislayer"] tspan,
:root .js-plotly-plot [class*="yaxis"] text,
:root .js-plotly-plot [class*="yaxis"] tspan,
:root .js-plotly-plot .updatemenu-container text,
:root .js-plotly-plot .updatemenu-container tspan { fill: #e8e8e8 !important; }

/* === Plot containers === */
.pc-plot { min-height: 600px !important; }

/* === Blocks / panels === */
.block, .panel, .form {
  background: var(--pc-surface) !important;
  border-color: var(--pc-border) !important;
}

/* === Dropdown items === */
.options, .option-list, ul[role="listbox"] {
  background: var(--pc-surface) !important;
  border-color: var(--pc-border) !important;
}
ul[role="listbox"] li,
.option { color: var(--pc-text) !important; }

/* === Chatbot (Gradio 6) — correct DOM structure ===
   In Gradio 6 the .chatbot class is on the INNER prose element inside each
   message bubble, NOT on the outer container.  That means ".chatbot .bot"
   is backwards — .bot is an ANCESTOR of .chatbot, not a descendant.
   Correct scoping: .message-wrap contains the .bot / .user bubble divs.
   .bubble-wrap is the outer scrollable pane. */
:root .bubble-wrap {
  background: var(--pc-surface) !important;
}
:root .message-wrap .bot {
  background-color: var(--pc-surface-alt) !important;
  border-color: var(--pc-border) !important;
}
:root .message-wrap .user {
  background-color: rgba(74,222,128,0.13) !important;
  border-color: rgba(74,222,128,0.4) !important;
}
:root .message-wrap .bot *,
:root .message-wrap .user * { color: var(--pc-text) !important; }

/* === Plotly charts — override SVG text in light mode === */
@media (prefers-color-scheme: light) {
  :root .js-plotly-plot [class*="tick"] text,
  :root .js-plotly-plot [class*="tick"] tspan,
  :root .js-plotly-plot [class*="axislayer"] text,
  :root .js-plotly-plot [class*="axislayer"] tspan,
  :root .js-plotly-plot [class*="yaxis"] text,
  :root .js-plotly-plot [class*="yaxis"] tspan,
  :root .js-plotly-plot .updatemenu-container text,
  :root .js-plotly-plot .updatemenu-container tspan { fill: #333333 !important; }
  :root .js-plotly-plot .modebar-btn path { fill: #555555 !important; }
}

/* === Team cards === */
.team-card {
  background: var(--pc-surface);
  border: 1px solid var(--pc-border);
  border-radius: 8px; padding: 10px; margin: 4px 0;
}
.teams-empty {
  color: var(--pc-text-muted);
  font-size: 13px; padding: 8px;
  border: 1px dashed var(--pc-border);
  border-radius: 6px; text-align: center;
}

/* === Banners === */
.date-ai-label { color: var(--pc-text-muted); font-size: 12px; font-style: italic; }
.banner-green {
  background: rgba(74,222,128,0.1);
  border: 1px solid var(--pc-green);
  border-radius: 8px; padding: 10px 14px;
  color: var(--pc-green); margin-bottom: 8px;
}
.banner-warning {
  background: rgba(251,191,36,0.1);
  border: 1px solid #fbbf24;
  border-radius: 8px; padding: 10px 14px;
  color: #b45309; margin-bottom: 8px;
}
@media (prefers-color-scheme: dark) {
  .banner-warning { color: #fbbf24; }
}

/* === Loading skeleton === */
.skeleton {
  background: linear-gradient(90deg,
    var(--pc-surface-alt) 25%, var(--pc-surface) 50%, var(--pc-surface-alt) 75%);
  background-size: 200% 100%;
  animation: shimmer 1.5s infinite;
  min-height: 300px; border-radius: 8px;
}
@keyframes shimmer {
  0%   { background-position: 200% 0; }
  100% { background-position: -200% 0; }
}
"""

SKELETON_HTML = '<div class="skeleton"></div>'


# ---------------------------------------------------------------------------
# Helper: teams display HTML
# ---------------------------------------------------------------------------
def _teams_html(teams: list) -> str:
    """Render team cards as colored pills/badges."""
    if not teams:
        return '<div class="teams-empty">No teams added yet</div>'
    parts = []
    for team in teams:
        color = team.get("color", "#4ade80")
        name = team.get("name", "")
        parts.append(
            '<div class="team-card" style="display:flex; align-items:center; gap:8px;">'
            f'<span style="display:inline-block; width:14px; height:14px; border-radius:50%; '
            f'background:{color}; flex-shrink:0;"></span>'
            f'<span style="font-size:14px;">{name}</span>'
            "</div>"
        )
    return "\n".join(parts)


# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
# Helper: extract team_name → color dict from state
# ---------------------------------------------------------------------------
def _state_team_colors(state: dict) -> dict[str, str]:
    """Return {team_name: color} from state inputs, for use in Gantt coloring."""
    teams = state.get("inputs", {}).get("teams", [])
    return {t["name"]: t.get("color", "#4ade80") for t in teams if isinstance(t, dict) and t.get("name")}


# ---------------------------------------------------------------------------
# Event handlers
# ---------------------------------------------------------------------------

def on_template_select(template_name, state):
    """Load template data and populate objective, tasks, teams in state."""
    import copy
    state = copy.deepcopy(state)

    if not template_name or template_name == "\u2014 None \u2014":
        state["inputs"]["template"] = None
        return (
            state["inputs"]["objective"],
            state["inputs"]["tasks"],
            _teams_html(state["inputs"]["teams"]),
            gr.update(choices=[t["name"] for t in state["inputs"]["teams"]]),
            state,
        )

    templates = _load_templates()
    template = templates.get(template_name)
    if template is None:
        return (
            state["inputs"]["objective"],
            state["inputs"]["tasks"],
            _teams_html(state["inputs"]["teams"]),
            gr.update(choices=[t["name"] for t in state["inputs"]["teams"]]),
            state,
        )

    state["inputs"]["template"] = template_name
    state["inputs"]["objective"] = template.get("objective", "")
    state["inputs"]["tasks"] = template.get("tasks", "")
    state["inputs"]["teams"] = template.get("teams", [])

    team_names = [t["name"] for t in state["inputs"]["teams"]]
    return (
        state["inputs"]["objective"],
        state["inputs"]["tasks"],
        _teams_html(state["inputs"]["teams"]),
        gr.update(choices=team_names),
        state,
    )


def on_add_team(name, color, state):
    """Sanitize name, append team dict to state, refresh display."""
    import copy
    state = copy.deepcopy(state)

    sanitized = name.strip() if name else ""
    if not sanitized:
        team_names = [t["name"] for t in state["inputs"]["teams"]]
        return (
            _teams_html(state["inputs"]["teams"]),
            gr.update(choices=team_names),
            "",
            "#4ade80",
            state,
        )

    team_id = sanitized.lower().replace(" ", "_")
    existing_ids = {t.get("id") for t in state["inputs"]["teams"]}
    if team_id not in existing_ids:
        state["inputs"]["teams"].append(
            {"id": team_id, "name": sanitized, "color": color or "#4ade80"}
        )

    team_names = [t["name"] for t in state["inputs"]["teams"]]
    return (
        _teams_html(state["inputs"]["teams"]),
        gr.update(choices=team_names),
        "",
        "#4ade80",
        state,
    )


def on_remove_team(team_name, state):
    """Filter out team with matching name, refresh display."""
    import copy
    state = copy.deepcopy(state)

    if team_name:
        state["inputs"]["teams"] = [
            t for t in state["inputs"]["teams"] if t["name"] != team_name
        ]

    team_names = [t["name"] for t in state["inputs"]["teams"]]
    return (
        _teams_html(state["inputs"]["teams"]),
        gr.update(choices=team_names, value=None),
        state,
    )


async def on_generate(objective, tasks, start_date, end_date, start_ai, end_ai, state):
    """Async generator: immediately yield skeletons, then generate and render the plan."""
    import copy
    print(f"[on_generate] CALLED — objective={objective!r:.60}")

    # --- Step 1: Skeleton / loading feedback ---
    # gr.Plot() only accepts figures or None — use None for plots during loading,
    # skeleton HTML only for the gr.HTML risks output.
    loading_html = SKELETON_HTML + '<div style="color:#a0a0a0;text-align:center;padding:8px;font-size:13px">Generating...</div>'
    yield (
        gr.update(value=None),          # gantt_plot  (gr.Plot — no strings)
        gr.update(value=None),          # raci_plot   (gr.Plot — no strings)
        gr.update(value=loading_html),  # risks_html_out (gr.HTML — strings OK)
        gr.update(value=None),          # heatmap_plot (gr.Plot — no strings)
        gr.update(value=""),            # date_banner
        gr.update(value=""),            # warning_banner
        gr.update(value='<div class="banner-green">Generating your project plan...</div>'),
        state,
    )

    # --- Step 2: Update state inputs ---
    state = copy.deepcopy(state)
    state["inputs"]["objective"] = objective or ""
    state["inputs"]["tasks"] = tasks or ""
    state["inputs"]["start_date"] = start_date or ""
    state["inputs"]["end_date"] = end_date or ""
    state["inputs"]["start_date_ai"] = bool(start_ai)
    state["inputs"]["end_date_ai"] = bool(end_ai)

    # Validate
    if not (objective or "").strip():
        yield (
            gr.update(value=""),
            gr.update(value=""),
            gr.update(value=""),
            gr.update(value=""),
            gr.update(value=""),
            gr.update(value=""),
            gr.update(
                value='<div class="banner-warning">'
                      "Please enter a project objective before generating."
                      "</div>"
            ),
            state,
        )
        return

    # --- Step 3: Call orchestrator ---
    try:
        inp = state["inputs"]
        result = await orchestrator.generate(
            objective=inp.get("objective", ""),
            tasks=inp.get("tasks", ""),
            teams=inp.get("teams", []),
            start_date=inp.get("start_date", ""),
            end_date=inp.get("end_date", ""),
            start_ai=inp.get("start_date_ai", False),
            end_ai=inp.get("end_date_ai", False),
        )
    except Exception as exc:
        yield (
            gr.update(value=None),  # gantt_plot
            gr.update(value=None),  # raci_plot
            gr.update(value='<div class="banner-warning">Generation failed.</div>'),
            gr.update(value=None),  # heatmap_plot
            gr.update(value=""),
            gr.update(value=""),
            gr.update(value=f'<div class="banner-warning">Generation failed: {exc}</div>'),
            state,
        )
        return

    # --- Step 4: Handle orchestrator error (treat empty string as error too) ---
    if isinstance(result, dict) and result.get("plan") is None:
        yield (
            gr.update(value=None),  # gantt_plot
            gr.update(value=None),  # raci_plot
            gr.update(value=f'<div class="banner-warning">{result["error"]}</div>'),
            gr.update(value=None),  # heatmap_plot
            gr.update(value=""),
            gr.update(value=""),
            gr.update(value=f'<div class="banner-warning">{result["error"]}</div>'),
            state,
        )
        return

    # --- Step 5: Persist plan in state (always store as dict) ---
    plan_raw = result.get("plan")
    print(f"[on_generate] result.error={result.get('error')!r}")
    print(f"[on_generate] plan_raw type={type(plan_raw).__name__}, "
          f"phases={len(plan_raw.get('phases', [])) if isinstance(plan_raw, dict) else 'N/A'}")
    plan_data = plan_raw or {}
    if not isinstance(plan_data, dict):
        plan_data = {}
    state["plan"] = plan_data
    state["original_plan"] = copy.deepcopy(plan_data)
    _team_colors = _state_team_colors(state)

    # --- Step 6: Build ProjectPlan and render charts ---
    try:
        project_plan = ProjectPlan.model_validate(plan_data)
        print(f"[on_generate] ProjectPlan built: {len(project_plan.phases)} phases, "
              f"{len(project_plan.raci)} raci, {len(project_plan.risks)} risks, "
              f"{len(project_plan.workload)} workload teams")
    except Exception as exc:
        print(f"[on_generate] ProjectPlan.model_validate FAILED: {exc}")
        project_plan = ProjectPlan()

    try:
        gantt_fig = gantt.render(project_plan, team_colors=_team_colors)
        print(f"[on_generate] gantt rendered OK")
    except Exception as exc:
        print(f"[on_generate] gantt FAILED: {exc}")
        gantt_fig = None

    try:
        raci_fig = raci.render(project_plan)
        print(f"[on_generate] raci rendered OK")
    except Exception as exc:
        print(f"[on_generate] raci FAILED: {exc}")
        raci_fig = None

    try:
        risks_html_content = risks.render(project_plan)
        print(f"[on_generate] risks rendered OK — {len(project_plan.risks)} risks in plan")
    except Exception as exc:
        print(f"[on_generate] risks FAILED: {exc}")
        risks_html_content = f'<div class="banner-warning">Could not render risks: {exc}</div>'

    try:
        heatmap_fig = heatmap.render(project_plan)
        print(f"[on_generate] heatmap rendered OK")
    except Exception as exc:
        print(f"[on_generate] heatmap FAILED: {exc}")
        heatmap_fig = None

    # --- Step 7: Banners ---
    date_banner_html = ""
    if plan_data.get("timeline_note"):
        date_banner_html = (
            f'<div class="banner-green">{plan_data["timeline_note"]}</div>'
        )

    warning_banner_html = ""
    if plan_data.get("timeline_warning"):
        warning_banner_html = (
            f'<div class="banner-warning">{plan_data["timeline_warning"]}</div>'
        )

    yield (
        gr.update(value=gantt_fig),
        gr.update(value=raci_fig),
        gr.update(value=risks_html_content),
        gr.update(value=heatmap_fig),
        gr.update(value=date_banner_html),
        gr.update(value=warning_banner_html),
        gr.update(value='<div class="banner-green">Plan generated successfully!</div>'),
        state,
    )


async def on_chat_submit(message, chat_history, state):
    """Handle chat refinement: call orchestrator.refine and re-render all charts."""
    import copy

    if not message or not message.strip():
        return chat_history, gr.update(), gr.update(), gr.update(), gr.update(), state

    plan_data = state.get("plan", {})
    chat_history = list(chat_history or [])

    if not plan_data or not plan_data.get("phases"):
        chat_history.append({"role": "user", "content": message})
        chat_history.append({"role": "assistant", "content": "No plan to refine yet. Please generate a plan first."})
        return chat_history, gr.update(), gr.update(), gr.update(), gr.update(), state

    try:
        current_plan = ProjectPlan.model_validate(plan_data)
        result = await orchestrator.refine(current_plan, message)
    except Exception as exc:
        chat_history.append({"role": "user", "content": message})
        chat_history.append({"role": "assistant", "content": f"Refinement failed: {exc}"})
        return chat_history, gr.update(), gr.update(), gr.update(), gr.update(), state

    if isinstance(result, dict) and result.get("error"):
        chat_history.append({"role": "user", "content": message})
        chat_history.append({"role": "assistant", "content": str(result["error"])})
        return chat_history, gr.update(), gr.update(), gr.update(), gr.update(), state

    # Update state with refined plan (always store as dict)
    state = copy.deepcopy(state)
    new_plan = result.get("plan") or {}
    if not isinstance(new_plan, dict):
        new_plan = {}
    state["plan"] = new_plan
    _team_colors = _state_team_colors(state)

    # Re-render charts
    try:
        project_plan = ProjectPlan.model_validate(new_plan)
    except Exception:
        project_plan = ProjectPlan()

    try:
        gantt_fig = gantt.render(project_plan, team_colors=_team_colors)
    except Exception:
        gantt_fig = None

    try:
        raci_fig = raci.render(project_plan)
    except Exception:
        raci_fig = None

    try:
        risks_html_content = risks.render(project_plan)
    except Exception as exc:
        risks_html_content = (
            f'<div class="banner-warning">Could not render risks: {exc}</div>'
        )

    try:
        heatmap_fig = heatmap.render(project_plan)
    except Exception:
        heatmap_fig = None

    chat_history.append({"role": "user", "content": message})
    chat_history.append({"role": "assistant", "content": "Plan updated! All charts have been refreshed."})

    return (
        chat_history,
        gr.update(value=gantt_fig),
        gr.update(value=raci_fig),
        gr.update(value=risks_html_content),
        gr.update(value=heatmap_fig),
        state,
    )


# ---------------------------------------------------------------------------
# AI checkbox toggle helpers
# ---------------------------------------------------------------------------
def toggle_start_date(start_ai):
    """Gray out start_date_input when AI is handling it."""
    return gr.update(interactive=not bool(start_ai))


def toggle_end_date(end_ai):
    """Gray out end_date_input when AI is handling it."""
    return gr.update(interactive=not bool(end_ai))


async def on_export_xlsx(state):
    """Export Gantt image, RACI matrix, and Risk Register to a single XLSX workbook."""
    plan_data = state.get("plan", {})
    if not plan_data or not plan_data.get("phases"):
        return None
    try:
        project_plan = ProjectPlan.model_validate(plan_data)
        path = await export_xlsx(project_plan, filename="projectcraft_plan")
        return path or None
    except Exception as exc:
        print(f"[export_xlsx] FAILED: {exc}")
        return None


# ---------------------------------------------------------------------------
# Gradio UI layout
# ---------------------------------------------------------------------------
with gr.Blocks() as demo:
    state = gr.State(value=INITIAL_STATE)

    with gr.Row():
        # ------------------------------------------------------------------ #
        # Left column — Input panel (~35%)
        # ------------------------------------------------------------------ #
        with gr.Column(scale=35):
            gr.Markdown("# \U0001f4ca ProjectCraft")

            # Template dropdown
            _template_names = list(_load_templates().keys())
            _template_choices = ["\u2014 None \u2014"] + _template_names
            template_dropdown = gr.Dropdown(
                label="Template",
                choices=_template_choices,
                value="\u2014 None \u2014",
            )

            # Objective
            objective_input = gr.Textbox(
                label="Project Objective",
                lines=4,
                placeholder="Describe your project goal...",
            )

            # Tasks / Milestones
            tasks_input = gr.Textbox(
                label="Key Tasks / Milestones",
                lines=3,
                placeholder="List the main tasks...",
            )

            # Teams section
            gr.Markdown("### \U0001f465 Teams")
            teams_display = gr.HTML(value=_teams_html([]))
            with gr.Row():
                team_name_input = gr.Textbox(
                    label="Team Name",
                    placeholder="e.g. Engineering",
                )
                team_color_input = gr.ColorPicker(label="Color", value="#4ade80")
            add_team_btn = gr.Button("+ Add Team", size="sm")
            remove_team_dropdown = gr.Dropdown(
                label="Remove team",
                choices=[],
                value=None,
            )
            remove_team_btn = gr.Button("Remove", size="sm")

            # Timeline section
            gr.Markdown("### \U0001f4c5 Timeline")
            start_date_input = gr.Textbox(label="Start Date (YYYY-MM-DD)")
            start_ai_checkbox = gr.Checkbox(label="Let AI decide start date")
            end_date_input = gr.Textbox(label="End Date (YYYY-MM-DD)")
            end_ai_checkbox = gr.Checkbox(label="Let AI decide end date")

            # Generate
            generate_btn = gr.Button("\U0001f680 Generate Plan", variant="primary")
            status_html = gr.HTML("")

        # ------------------------------------------------------------------ #
        # Right column — Output panel (~65%)
        # ------------------------------------------------------------------ #
        with gr.Column(scale=65):
            # Date / warning banners
            date_banner = gr.HTML("")
            warning_banner = gr.HTML("")

            # Four output tabs
            with gr.Tabs():
                with gr.Tab("\U0001f4ca Gantt"):
                    gantt_plot = gr.Plot(visible=True, show_label=False, elem_classes=["pc-plot"])
                with gr.Tab("\U0001f465 RACI"):
                    raci_plot = gr.Plot(visible=True, show_label=False, elem_classes=["pc-plot"])
                with gr.Tab("\u26a0\ufe0f Risks"):
                    risks_html_out = gr.HTML(visible=True)
                with gr.Tab("\U0001f525 Workload"):
                    heatmap_plot = gr.Plot(visible=True, show_label=False, elem_classes=["pc-plot"])

            # Export
            export_xlsx_btn = gr.DownloadButton("📥 Export XLSX", size="sm")

            # Chat refinement
            gr.Markdown("### \U0001f4ac Refine Plan")
            chatbot = gr.Chatbot(height=160, show_label=False)
            chat_input = gr.Textbox(
                label="",
                placeholder="e.g. Add a testing phase after development",
            )
            chat_submit_btn = gr.Button("Send", variant="primary")


    # ---------------------------------------------------------------------- #
    # Event wiring
    # ---------------------------------------------------------------------- #

    # AI date toggles — gray out the text fields when checked
    start_ai_checkbox.change(
        fn=toggle_start_date,
        inputs=[start_ai_checkbox],
        outputs=[start_date_input],
    )
    end_ai_checkbox.change(
        fn=toggle_end_date,
        inputs=[end_ai_checkbox],
        outputs=[end_date_input],
    )

    # Template selection — populate fields
    template_dropdown.change(
        fn=on_template_select,
        inputs=[template_dropdown, state],
        outputs=[
            objective_input,
            tasks_input,
            teams_display,
            remove_team_dropdown,
            state,
        ],
    )

    # Add team
    add_team_btn.click(
        fn=on_add_team,
        inputs=[team_name_input, team_color_input, state],
        outputs=[
            teams_display,
            remove_team_dropdown,
            team_name_input,
            team_color_input,
            state,
        ],
    )

    # Remove team
    remove_team_btn.click(
        fn=on_remove_team,
        inputs=[remove_team_dropdown, state],
        outputs=[teams_display, remove_team_dropdown, state],
    )

    # Generate plan — async generator (streaming with yield)
    generate_btn.click(
        fn=on_generate,
        inputs=[
            objective_input,
            tasks_input,
            start_date_input,
            end_date_input,
            start_ai_checkbox,
            end_ai_checkbox,
            state,
        ],
        outputs=[
            gantt_plot,
            raci_plot,
            risks_html_out,
            heatmap_plot,
            date_banner,
            warning_banner,
            status_html,
            state,
        ],
    )

    # Chat refinement — wired to both button and Enter key
    _chat_inputs = [chat_input, chatbot, state]
    _chat_outputs = [chatbot, gantt_plot, raci_plot, risks_html_out, heatmap_plot, state]

    chat_submit_btn.click(
        fn=on_chat_submit,
        inputs=_chat_inputs,
        outputs=_chat_outputs,
    )
    chat_input.submit(
        fn=on_chat_submit,
        inputs=_chat_inputs,
        outputs=_chat_outputs,
    )

    # XLSX export
    export_xlsx_btn.click(
        fn=on_export_xlsx,
        inputs=[state],
        outputs=[export_xlsx_btn],
    )


# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, css=CSS, ssr_mode=False)