from __future__ import annotations

import os
import shutil
from pathlib import Path

from django.core.management.base import BaseCommand

from control_plane.resources.models import Resource, ResourceKind


class Command(BaseCommand):
    help = "Seed DGX Spark compute and model resources for Sol and Qwen."

    def opencode_command(self, model_key: str) -> str:
        env_name = f"ARTIFEX_{model_key.upper()}_OPENCODE_COMMAND"
        if os.environ.get(env_name):
            return os.environ[env_name]
        if os.environ.get("ARTIFEX_OPENCODE_COMMAND"):
            return os.environ["ARTIFEX_OPENCODE_COMMAND"]
        executable = shutil.which("opencode")
        if executable is None:
            default_user_install = Path.home() / ".opencode" / "bin" / "opencode"
            executable = str(default_user_install) if default_user_install.exists() else "opencode"
        return f"{executable} run"

    def opencode_config(self, model_key: str) -> dict[str, object]:
        return {
            "model_key": model_key,
            "transport": os.environ.get(f"ARTIFEX_{model_key.upper()}_TRANSPORT", os.environ.get("ARTIFEX_OPENCODE_TRANSPORT", "local")),
            "command": self.opencode_command(model_key),
            "use_pty": os.environ.get(f"ARTIFEX_{model_key.upper()}_USE_PTY", os.environ.get("ARTIFEX_OPENCODE_USE_PTY", "0")) == "1",
            "timeout_seconds": int(os.environ.get(f"ARTIFEX_{model_key.upper()}_TIMEOUT_SECONDS", os.environ.get("ARTIFEX_OPENCODE_TIMEOUT_SECONDS", "120"))),
        }

    def update_opencode_resource(self, *, model_key: str, name: str, roles: list[str], compute: Resource) -> None:
        defaults = {
            "name": name,
            "kind": ResourceKind.MODEL,
            "provider": "opencode",
            "compute": compute,
            "roles": roles,
            "config": self.opencode_config(model_key),
        }
        existing = next((resource for resource in Resource.objects.filter(provider="opencode") if resource.config.get("model_key") == model_key), None)
        if existing is not None:
            for field, value in defaults.items():
                setattr(existing, field, value)
            existing.save(update_fields=[*defaults.keys(), "updated_at"])
            return
        Resource.objects.update_or_create(name=name, defaults={key: value for key, value in defaults.items() if key != "name"})

    def handle(self, *args, **options):
        spark, _ = Resource.objects.update_or_create(
            name=os.environ.get("ARTIFEX_SPARK_RESOURCE_NAME", "DGX Spark"),
            defaults={
                "kind": ResourceKind.MODEL_COMPUTE,
                "provider": "ssh",
                "config": {"ssh_alias": os.environ.get("ARTIFEX_SPARK_SSH_ALIAS", "spark")},
                "roles": [],
            },
        )
        self.update_opencode_resource(
            model_key="sol",
            name=os.environ.get("ARTIFEX_SOL_RESOURCE_NAME", "GPT-5.6 Sol"),
            roles=["PROJECT_BRAIN", "PLANNING", "ARCHAEOLOGY_INTERPRETATION", "AGENT_DESIGN", "ESCALATION"],
            compute=spark,
        )
        self.update_opencode_resource(
            model_key="terra",
            name=os.environ.get("ARTIFEX_TERRA_RESOURCE_NAME", "GPT-5.6 Terra"),
            roles=["REASONING", "REVIEW", "PORTFOLIO_IC", "STRATEGY"],
            compute=spark,
        )
        self.update_opencode_resource(
            model_key="luna",
            name=os.environ.get("ARTIFEX_LUNA_RESOURCE_NAME", "GPT-5.6 Luna"),
            roles=["REASONING", "MARKET_RESEARCH", "WEB_RESEARCH", "SYNTHESIS"],
            compute=spark,
        )
        Resource.objects.update_or_create(
            name=os.environ.get("ARTIFEX_QWEN_RESOURCE_NAME", "Qwen"),
            defaults={
                "kind": ResourceKind.MODEL,
                "provider": "local_inference",
                "compute": spark,
                "roles": ["CODING", "REVIEW", "REASONING"],
                "config": {
                    "endpoint_url": os.environ.get("ARTIFEX_QWEN_ENDPOINT_URL", "http://192.168.1.162:8002/v1/chat/completions"),
                    "health_url": os.environ.get("ARTIFEX_QWEN_HEALTH_URL", "http://192.168.1.162:8002/health"),
                    "model": os.environ.get("ARTIFEX_QWEN_MODEL", "qwen38"),
                    "extra_body": {"chat_template_kwargs": {"enable_thinking": False}},
                },
            },
        )
        Resource.objects.update_or_create(
            name=os.environ.get("ARTIFEX_SEARXNG_RESOURCE_NAME", "SearXNG Search"),
            defaults={
                "kind": ResourceKind.COMPUTE,
                "provider": "searxng",
                "compute": spark,
                "roles": ["WEB_RESEARCH", "MARKET_RESEARCH", "SEARCH"],
                "config": {
                    "endpoint_url": os.environ.get("ARTIFEX_SEARXNG_ENDPOINT_URL", "http://127.0.0.1:8080"),
                    "timeout_seconds": int(os.environ.get("ARTIFEX_SEARXNG_TIMEOUT_SECONDS", "20")),
                },
            },
        )
        self.stdout.write(self.style.SUCCESS("Seeded Spark model resources."))