"""
SOG-X Modular Studio - Dynamic Swagger-Style Interface
Auto-discovers backend routes via OpenAPI and mirrors Swagger UI's own behavior:
- multipart/form-data bodies -> one widget per field (including real file upload)
- application/json bodies    -> a single JSON editor, pre-filled with an example
                                 built from each field's default/example value
- query/path parameters      -> individual widgets, same as Swagger's parameter list
Compatible with Gradio 6.0+
"""
import json
import os
import gradio as gr
import spaces
from client import BackendClient

BACKEND_URL = os.getenv("BACKEND_URL", "http://localhost:7860")
HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("API_KEY", "")
USER_KEY = os.getenv("USER_KEY", "")  # never hardcode a fallback secret here

client = BackendClient(backend_url=BACKEND_URL, token=HF_TOKEN)

FILE_NAME_HINTS = ("file", "document", "upload", "attachment", "audio", "video", "image")


@spaces.GPU
def _warmup():
    """Dummy function so HF's ZeroGPU startup check finds a @spaces.GPU function."""
    return "ok"


custom_css = """
.module-radio label {
    background-color: #2b2b3b !important; border: 1px solid #3f3f56 !important;
    border-radius: 8px !important; margin-bottom: 6px !important; padding: 10px 14px !important;
    transition: all 0.2s ease;
}
.module-radio label:hover { border-color: #6c63ff !important; }
"""


# ── Schema Loading ──────────────────────────────────────────────────────
def load_backend_modules(api_key: str):
    try:
        ok, schema = client.fetch_openapi_schema(user_api_key=api_key)
    except Exception as e:
        ok, schema = False, {"error": "NETWORK_ERROR", "detail": str(e)}

    modules_map, error_detail = {}, None

    if not ok:
        error_detail = f"{schema.get('error', 'UNKNOWN')}: {schema.get('detail', '')}".strip(": ")
    elif isinstance(schema, dict) and "paths" not in schema:
        error_detail = "Schema fetched but has no 'paths' key — check the response shape."

    if ok and isinstance(schema, dict) and "paths" in schema:
        for path, methods in schema["paths"].items():
            for method, details in methods.items():
                if method.lower() not in ["get", "post"]:
                    continue
                tags = details.get("tags", ["General"])
                tag_name = tags[0].capitalize() if tags else "General"
                summary = details.get("summary") or f"{method.upper()} {path}"
                modules_map[f"[{tag_name}] {summary}"] = {
                    "path": path,
                    "method": method.upper(),
                    "summary": summary,
                    "description": details.get("description", "No description provided."),
                    "parameters": details.get("parameters", []),
                    "request_body": details.get("requestBody", {}),
                }

    if modules_map:
        module_names = list(modules_map.keys())
    elif error_detail:
        module_names = [f"⚠️ Discovery failed: {error_detail}"]
    else:
        module_names = ["No Modules Discovered"]

    full_schema = schema if (ok and isinstance(schema, dict)) else {}
    return modules_map, module_names, full_schema


def resolve_ref(full_schema: dict, node):
    """Resolve a single-level '$ref' like '#/components/schemas/ChatRequest'."""
    if not isinstance(node, dict):
        return {}
    if "$ref" in node:
        parts = node["$ref"].lstrip("#/").split("/")
        resolved = full_schema
        for p in parts:
            resolved = resolved.get(p, {}) if isinstance(resolved, dict) else {}
        return resolved
    return node


def _is_file_schema(name: str, schema: dict, is_multipart: bool) -> bool:
    if schema.get("format") == "binary" or "contentMediaType" in schema:
        return True
    if is_multipart and name.lower() in FILE_NAME_HINTS:
        return True
    return False


def _field_from_schema(name, in_, required, description, schema, full_schema, is_multipart=False):
    schema = resolve_ref(full_schema, schema)
    items = resolve_ref(full_schema, schema.get("items")) if schema.get("items") else {}
    return {
        "name": name,
        "in": in_,
        "required": required,
        "type": schema.get("type", "string"),
        "is_file": _is_file_schema(name, schema, is_multipart),
        "enum": schema.get("enum"),
        "items": items,
        "pattern": schema.get("pattern"),
        "default": schema.get("default"),
        "example": schema.get("example"),
        "description": description or schema.get("description", ""),
    }


def extract_param_fields(mod: dict, full_schema: dict):
    """Query / path parameters — Swagger always renders these as individual widgets."""
    fields, seen = [], set()
    for p in mod.get("parameters", []):
        name = p.get("name")
        if not name or name in seen:
            continue
        seen.add(name)
        fields.append(_field_from_schema(
            name, p.get("in", "query"), bool(p.get("required", False)),
            p.get("description", ""), p.get("schema", {}) or {}, full_schema,
        ))
    return fields


def get_request_body(mod: dict, full_schema: dict):
    """Returns (content_type, resolved_body_schema) or (None, None) if there's no body."""
    rb = mod.get("request_body") or {}
    content = rb.get("content") or {}
    for ctype, cinfo in content.items():
        return ctype, resolve_ref(full_schema, cinfo.get("schema"))
    return None, None


def extract_multipart_fields(body_schema: dict, full_schema: dict):
    fields, required_names = [], set(body_schema.get("required") or [])
    for name, prop in (body_schema.get("properties") or {}).items():
        fields.append(_field_from_schema(
            name, "body", name in required_names,
            prop.get("description", ""), prop, full_schema, is_multipart=True,
        ))
    return fields


def _placeholder_for_type(t: str):
    return {"string": "", "integer": 0, "number": 0, "boolean": False, "array": [], "object": {}}.get(t, "")


def build_json_example(body_schema: dict, full_schema: dict) -> str:
    """Mirrors Swagger's 'Example Value' tab: use an explicit example if given,
    else synthesize one from each field's own default/example, falling back to a
    type-shaped placeholder."""
    if "example" in body_schema:
        example = body_schema["example"]
    elif body_schema.get("examples"):
        ex = body_schema["examples"]
        example = list(ex.values())[0] if isinstance(ex, dict) else ex[0]
    else:
        example = {}
        for name, prop in (body_schema.get("properties") or {}).items():
            prop = resolve_ref(full_schema, prop)
            if "example" in prop:
                example[name] = prop["example"]
            elif "default" in prop:
                example[name] = prop["default"]
            else:
                example[name] = _placeholder_for_type(prop.get("type", "string"))
    return json.dumps(example, indent=2)


def build_field_component(f: dict, key: str):
    """Pick the widget type the way Swagger UI would, based on the field's schema."""
    label = f["name"] + ("  *required" if f["required"] else "")
    default = f["example"] if f["example"] is not None else f["default"]

    if f["is_file"]:
        return gr.File(label=label, key=key)
    if f["enum"]:
        return gr.Dropdown(label=label, choices=[str(v) for v in f["enum"]],
                            value=str(default) if default is not None else None,
                            info=f["description"], interactive=True, key=key)
    if f["type"] == "boolean":
        return gr.Checkbox(label=label, value=bool(default) if default is not None else False,
                            info=f["description"], interactive=True, key=key)
    if f["type"] in ("integer", "number"):
        return gr.Number(label=label, value=default, info=f["description"], interactive=True, key=key)
    if f["type"] == "array":
        item_enum = f["items"].get("enum") if isinstance(f["items"], dict) else None
        if item_enum:
            return gr.Dropdown(label=label, choices=[str(v) for v in item_enum], multiselect=True,
                                info=f["description"], interactive=True, key=key)
        return gr.Textbox(label=label, value=", ".join(default) if isinstance(default, list) else "",
                           placeholder="Comma-separated values", info=f["description"] or "List — separate items with commas",
                           interactive=True, key=key)

    placeholder = f["description"] or f["name"]
    if f["pattern"] and default is None:
        placeholder = f"Pattern: {f['pattern']}"
    return gr.Textbox(label=label, value=str(default) if default is not None else "", placeholder=placeholder,
                       interactive=True, key=key)


def coerce_field_value(f: dict, val):
    if val in (None, ""):
        return None
    if f["type"] == "array":
        return val if isinstance(val, list) else [v.strip() for v in str(val).split(",") if v.strip()]
    return val


INITIAL_MAP, INITIAL_CHOICES, INITIAL_SCHEMA = load_backend_modules(USER_KEY)

with gr.Blocks(title="SOG-X Modular Workspace") as demo:
    schema_state = gr.State(value=INITIAL_MAP)
    full_schema_state = gr.State(value=INITIAL_SCHEMA)

    with gr.Sidebar(position="left", open=True, label="SOG-X Studio") as studio_sidebar:
        gr.Markdown("## 🤖 SOG-X Studio")
        with gr.Accordion("⚙️ Backend & Credentials", open=False):
            gr.Markdown(f"**Backend URL:** `{BACKEND_URL}`")
            user_key_input = gr.Textbox(label="Active API Key", value="", type="password",
                                         placeholder="Enter your API key...", interactive=True)
            sync_btn = gr.Button("🔄 Sync Backend Modules", size="sm")

        gr.Markdown("---")
        gr.Markdown("### 📂 Active Modules")
        sidebar_menu = gr.Radio(choices=INITIAL_CHOICES,
                                 value=INITIAL_CHOICES[0] if INITIAL_CHOICES else None,
                                 label="Select Module", interactive=True, elem_classes=["module-radio"])

    @gr.render(inputs=[sidebar_menu, schema_state, full_schema_state])
    def render_workspace(selected_name, current_map, full_schema):
        if not selected_name or selected_name not in current_map:
            gr.Markdown("## Select a module from the sidebar")
            gr.Markdown("Choose any active backend route from the left navigation panel.")
            return

        mod = current_map[selected_name]
        gr.Markdown(f"## `{mod['method']}` {mod['path']}")
        gr.Markdown(f"**Summary:** {mod['summary']}\n\n{mod['description']}")
        gr.Markdown("---")
        gr.Markdown("### 📥 Input")

        param_fields = extract_param_fields(mod, full_schema)
        param_components = {f["name"]: build_field_component(f, key=f"param::{f['name']}") for f in param_fields}

        content_type, body_schema = get_request_body(mod, full_schema)
        body_mode, body_fields, body_components, json_editor = None, [], {}, None

        if content_type == "multipart/form-data" and body_schema:
            body_mode = "form"
            body_fields = extract_multipart_fields(body_schema, full_schema)
            body_components = {f["name"]: build_field_component(f, key=f"body::{f['name']}") for f in body_fields}
        elif content_type == "application/json" and body_schema:
            body_mode = "json"
            example_str = build_json_example(body_schema, full_schema)
            json_editor = gr.Code(label="Request Body (JSON) — edit before running",
                                   language="json", value=example_str, lines=10, interactive=True, key="json_editor")

        if not param_fields and not body_mode:
            gr.Markdown("_This endpoint takes no parameters — just click Run._")

        run_btn = gr.Button("🚀 Run", variant="primary", size="lg", key="run_btn")
        gr.Markdown("---")
        gr.Markdown("### 📤 Output")
        out_json = gr.JSON(label="Response Payload", visible=True, key="out_json")
        out_image = gr.Image(label="Response Image", visible=False, key="out_image")
        out_audio = gr.Audio(label="Response Audio", visible=False, key="out_audio")
        out_video = gr.Video(label="Response Video", visible=False, key="out_video")
        out_file = gr.File(label="Response File", visible=False, key="out_file")
        outputs = [out_json, out_image, out_audio, out_video, out_file]

        def _hide_all_but(target, value):
            result = {o: gr.update(visible=False) for o in outputs if o is not target}
            result[target] = gr.update(value=value, visible=True)
            return result

        run_inputs = [user_key_input] + list(param_components.values())
        if body_mode == "form":
            run_inputs += list(body_components.values())
        elif body_mode == "json":
            run_inputs += [json_editor]

        def _run(active_key, *values):
            n_params = len(param_fields)
            param_values = values[:n_params]
            rest = values[n_params:]

            resolved_path = mod["path"]
            params = {}
            for f, raw in zip(param_fields, param_values):
                val = coerce_field_value(f, raw)
                if val is None:
                    continue
                if f["in"] == "path":
                    resolved_path = resolved_path.replace("{" + f["name"] + "}", str(val))
                else:
                    params[f["name"]] = val

            payload, files, file_handles = {}, {}, []
            try:
                if body_mode == "json":
                    raw_json = rest[0] if rest else "{}"
                    try:
                        payload = json.loads(raw_json) if raw_json.strip() else {}
                    except Exception as e:
                        return _hide_all_but(out_json, {"error": f"Invalid JSON in request body: {e}"})
                elif body_mode == "form":
                    for f, raw in zip(body_fields, rest):
                        val = coerce_field_value(f, raw)
                        if val is None:
                            continue
                        if f["is_file"]:
                            file_path = val.name if hasattr(val, "name") else str(val)
                            fh = open(file_path, "rb")
                            file_handles.append(fh)
                            files[f["name"]] = fh
                        else:
                            payload[f["name"]] = val

                success, response = client.execute_dynamic_request(
                    path=resolved_path, method=mod["method"],
                    params=params, payload=payload, files=files,
                    user_api_key=active_key,
                )
            finally:
                for fh in file_handles:
                    fh.close()

            if not success:
                return _hide_all_but(out_json, {"error": str(response)})

            if isinstance(response, dict) and response.get("__binary__"):
                ctype, path = response.get("content_type", ""), response.get("file_path")
                if ctype.startswith("image/"):
                    return _hide_all_but(out_image, path)
                if ctype.startswith("audio/"):
                    return _hide_all_but(out_audio, path)
                if ctype.startswith("video/"):
                    return _hide_all_but(out_video, path)
                return _hide_all_but(out_file, path)

            result = _hide_all_but(out_json, response)
            if isinstance(response, dict) and "file_path" in response:
                result[out_file] = gr.update(value=response["file_path"], visible=True)
            return result

        run_btn.click(fn=_run, inputs=run_inputs, outputs=outputs)

    def refresh_schema_handler(active_key):
        new_map, new_choices, new_schema = load_backend_modules(active_key)
        selected = new_choices[0] if new_choices else None
        return new_map, new_schema, gr.Radio(choices=new_choices, value=selected)

    sync_btn.click(fn=refresh_schema_handler, inputs=[user_key_input],
                   outputs=[schema_state, full_schema_state, sidebar_menu])

    # Auto-collapse the sidebar the moment a module is picked, so the workspace
    # takes over the screen — same pattern as a mobile nav drawer closing on tap.
    sidebar_menu.change(fn=lambda: gr.Sidebar(open=False), outputs=studio_sidebar)

if __name__ == "__main__":
    _warmup()
    port = int(os.getenv("PORT", 7860))
    demo.launch(server_name="0.0.0.0", server_port=port,
                theme=gr.themes.Soft(), css=custom_css, share=False)