from __future__ import annotations

import time
import uuid
from enum import Enum
from pathlib import Path
from typing import Any

import structlog
from fastapi import APIRouter, Body, HTTPException
from pydantic import BaseModel, Field

from proxym.storage import DEFAULT_SQLITE_PATH, SQLiteNamespaceStore
from proxym.tools import ToolRegistry

logger = structlog.get_logger()
router = APIRouter(tags=["environments"])

_ENV_NAMESPACE = "environments"
_ENV_ASSIGNMENT_NAMESPACE = "environment_tool_assignments"


class EnvironmentKind(str, Enum):
    LOCAL_PATH = "local_path"
    SSH = "ssh"
    VM_LOCAL = "vm_local"
    PROXMOX_REMOTE = "proxmox_remote"
    RDP = "rdp"


class Environment(BaseModel):
    id: str = Field(default_factory=lambda: f"env-{uuid.uuid4().hex[:8]}")
    name: str
    kind: EnvironmentKind = EnvironmentKind.LOCAL_PATH
    path: str = ""
    host: str = ""
    user: str = ""
    port: int = 22
    vm_name: str = ""
    proxmox_node: str = ""
    proxmox_vm_id: str = ""
    notes: str = ""
    active: bool = True
    created_at: float = Field(default_factory=time.time)
    updated_at: float = Field(default_factory=time.time)

    def _ssh_target_label(self) -> str:
        prefix = f"{self.user}@" if self.user else ""
        port = f":{self.port}" if self.port else ""
        tail = f" {self.path}" if self.path else ""
        return f"ssh://{prefix}{self.host}{port}{tail}".strip()

    def _vm_target_label(self) -> str:
        parts = [self.vm_name or "local VM"]
        if self.path:
            parts.append(self.path)
        return " · ".join(parts)

    def _proxmox_target_label(self) -> str:
        prefix = f"{self.user}@" if self.user else ""
        port = f":{self.port}" if self.port else ""
        details = []
        if self.proxmox_node:
            details.append(f"node={self.proxmox_node}")
        if self.proxmox_vm_id:
            details.append(f"vm={self.proxmox_vm_id}")
        if self.path:
            details.append(self.path)
        suffix = f" {' '.join(details)}" if details else ""
        return f"proxmox://{prefix}{self.host}{port}{suffix}".strip()

    def _rdp_target_label(self) -> str:
        prefix = f"{self.user}@" if self.user else ""
        port = f":{self.port}" if self.port else ""
        tail = f" {self.path}" if self.path else ""
        return f"rdp://{prefix}{self.host}{port}{tail}".strip()

    def target_label(self) -> str:
        label_builders = {
            EnvironmentKind.LOCAL_PATH: lambda: self.path or "local path",
            EnvironmentKind.SSH: self._ssh_target_label,
            EnvironmentKind.VM_LOCAL: self._vm_target_label,
            EnvironmentKind.PROXMOX_REMOTE: self._proxmox_target_label,
            EnvironmentKind.RDP: self._rdp_target_label,
        }
        builder = label_builders.get(self.kind)
        return builder() if builder else self.name

    def summary(self, tools: list[str] | None = None) -> dict[str, Any]:
        assigned_tools = tools or []
        return {
            "id": self.id,
            "name": self.name,
            "kind": self.kind.value,
            "path": self.path,
            "host": self.host,
            "user": self.user,
            "port": self.port,
            "vm_name": self.vm_name,
            "proxmox_node": self.proxmox_node,
            "proxmox_vm_id": self.proxmox_vm_id,
            "notes": self.notes,
            "active": self.active,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "target": self.target_label(),
            "tools": assigned_tools,
            "tool_count": len(assigned_tools),
        }


class EnvironmentCreateRequest(BaseModel):
    name: str
    kind: EnvironmentKind = EnvironmentKind.LOCAL_PATH
    path: str = ""
    host: str = ""
    user: str = ""
    port: int = 22
    vm_name: str = ""
    proxmox_node: str = ""
    proxmox_vm_id: str = ""
    notes: str = ""
    active: bool = True


class EnvironmentUpdateRequest(BaseModel):
    name: str | None = None
    kind: EnvironmentKind | None = None
    path: str | None = None
    host: str | None = None
    user: str | None = None
    port: int | None = None
    vm_name: str | None = None
    proxmox_node: str | None = None
    proxmox_vm_id: str | None = None
    notes: str | None = None
    active: bool | None = None


class ToolEnvironmentAssignRequest(BaseModel):
    tool: str = Field(..., description="Tool name")
    environment_id: str = Field("", description="Environment ID (empty to unassign)")


class EnvironmentRegistry:
    def __init__(self, persist_path: Path | str | None = None) -> None:
        self._persist_path = Path(persist_path).expanduser() if persist_path is not None else None
        self._store = SQLiteNamespaceStore(self._persist_path) if self._persist_path is not None else None
        self.environments: dict[str, Environment] = {}
        self.tool_assignments: dict[str, str] = {}
        if self._persist_path and self._persist_path.exists():
            self._load()

    def _load(self) -> None:
        if self._store is None:
            return
        environment_payloads = self._store.load_namespace(_ENV_NAMESPACE)
        self.environments = {
            env_id: Environment(**payload)
            for env_id, payload in environment_payloads.items()
        }
        assignment_payloads = self._store.load_namespace(_ENV_ASSIGNMENT_NAMESPACE)
        self.tool_assignments = {
            tool_name: str(payload.get("environment_id", ""))
            for tool_name, payload in assignment_payloads.items()
            if payload.get("environment_id")
        }

    def _save(self) -> None:
        if self._persist_path is None:
            return
        if self._store is None:
            self._store = SQLiteNamespaceStore(self._persist_path)
        self._store.save_namespaces({
            _ENV_NAMESPACE: {
                env_id: env.model_dump(mode="json")
                for env_id, env in self.environments.items()
            },
            _ENV_ASSIGNMENT_NAMESPACE: {
                tool_name: {"tool": tool_name, "environment_id": environment_id}
                for tool_name, environment_id in self.tool_assignments.items()
                if environment_id
            },
        })

    def _tools_by_environment(self) -> dict[str, list[str]]:
        mapping: dict[str, list[str]] = {env_id: [] for env_id in self.environments}
        for tool_name, environment_id in self.tool_assignments.items():
            if environment_id in mapping:
                mapping[environment_id].append(tool_name)
        for tool_names in mapping.values():
            tool_names.sort()
        return mapping

    def list_environments(self) -> list[Environment]:
        return sorted(self.environments.values(), key=lambda env: env.name.lower())

    def list_payloads(self) -> list[dict[str, Any]]:
        tools_by_environment = self._tools_by_environment()
        return [
            environment.summary(tools_by_environment.get(environment.id, []))
            for environment in self.list_environments()
        ]

    def get(self, environment_id: str) -> Environment | None:
        return self.environments.get(environment_id)

    def create(self, environment: Environment) -> Environment:
        self.environments[environment.id] = environment
        self._save()
        return environment

    def update(self, environment_id: str, updates: dict[str, Any]) -> Environment | None:
        environment = self.environments.get(environment_id)
        if not environment:
            return None
        for key, value in updates.items():
            if hasattr(environment, key) and value is not None:
                setattr(environment, key, value)
        environment.updated_at = time.time()
        self._save()
        return environment

    def delete(self, environment_id: str) -> bool:
        if environment_id not in self.environments:
            return False
        del self.environments[environment_id]
        self.tool_assignments = {
            tool_name: assigned_environment_id
            for tool_name, assigned_environment_id in self.tool_assignments.items()
            if assigned_environment_id != environment_id
        }
        self._save()
        return True

    def assign_tool(self, tool_name: str, environment_id: str | None) -> Environment | None:
        if environment_id:
            if environment_id not in self.environments:
                return None
            self.tool_assignments[tool_name] = environment_id
            self._save()
            return self.environments[environment_id]
        self.tool_assignments.pop(tool_name, None)
        self._save()
        return None

    def get_tool_environment_id(self, tool_name: str) -> str:
        return self.tool_assignments.get(tool_name, "")

    def get_tool_environment(self, tool_name: str) -> Environment | None:
        environment_id = self.get_tool_environment_id(tool_name)
        return self.environments.get(environment_id) if environment_id else None

    def summary(self) -> dict[str, Any]:
        by_kind: dict[str, int] = {}
        for environment in self.environments.values():
            by_kind[environment.kind.value] = by_kind.get(environment.kind.value, 0) + 1
        return {
            "total": len(self.environments),
            "assigned_tools": len(self.tool_assignments),
            "by_kind": by_kind,
        }


_registry: EnvironmentRegistry | None = None


def get_environment_registry() -> EnvironmentRegistry:
    global _registry
    if _registry is None:
        _registry = EnvironmentRegistry(DEFAULT_SQLITE_PATH)
    return _registry


def get_tool_environment_snapshot(tool_name: str) -> dict[str, Any] | None:
    environment = get_environment_registry().get_tool_environment(tool_name)
    if not environment:
        return None
    registry = get_environment_registry()
    tools_by_environment = registry._tools_by_environment()
    return environment.summary(tools_by_environment.get(environment.id, []))


@router.get("/environments")
async def list_environments() -> dict[str, Any]:
    registry = get_environment_registry()
    return {
        "environments": registry.list_payloads(),
        "tool_assignments": registry.tool_assignments,
        "summary": registry.summary(),
    }


@router.get("/environments/assignments")
async def list_environment_assignments() -> dict[str, Any]:
    registry = get_environment_registry()
    tools_by_environment = registry._tools_by_environment()
    assignments = []
    for tool_name, environment_id in sorted(registry.tool_assignments.items()):
        environment = registry.get(environment_id)
        assignments.append({
            "tool": tool_name,
            "environment_id": environment_id,
            "environment": environment.summary(tools_by_environment.get(environment.id, [])) if environment else None,
        })
    return {
        "tool_assignments": registry.tool_assignments,
        "assignments": assignments,
        "summary": registry.summary(),
    }


@router.get("/environments/{environment_id}")
async def get_environment(environment_id: str) -> dict[str, Any]:
    registry = get_environment_registry()
    environment = registry.get(environment_id)
    if not environment:
        raise HTTPException(404, f"Environment '{environment_id}' not found")
    tools_by_environment = registry._tools_by_environment()
    return environment.summary(tools_by_environment.get(environment.id, []))


@router.post("/environments")
async def create_environment(req: EnvironmentCreateRequest) -> dict[str, Any]:
    registry = get_environment_registry()
    environment = Environment(**req.model_dump())
    created = registry.create(environment)
    logger.info("environment_created", environment=created.id, kind=created.kind.value)
    return created.summary()


@router.put("/environments/{environment_id}")
async def update_environment(environment_id: str, req: EnvironmentUpdateRequest) -> dict[str, Any]:
    registry = get_environment_registry()
    updates = {key: value for key, value in req.model_dump().items() if value is not None}
    environment = registry.update(environment_id, updates)
    if not environment:
        raise HTTPException(404, f"Environment '{environment_id}' not found")
    logger.info("environment_updated", environment=environment_id)
    tools_by_environment = registry._tools_by_environment()
    return environment.summary(tools_by_environment.get(environment.id, []))


@router.delete("/environments/{environment_id}")
async def delete_environment(environment_id: str) -> dict[str, Any]:
    registry = get_environment_registry()
    if not registry.delete(environment_id):
        raise HTTPException(404, f"Environment '{environment_id}' not found")
    logger.info("environment_deleted", environment=environment_id)
    return {"ok": True, "environment_id": environment_id}


@router.post("/environments/assign")
async def assign_tool_to_environment(req: ToolEnvironmentAssignRequest) -> dict[str, Any]:
    tool = ToolRegistry.get().get_tool(req.tool)
    if not tool:
        raise HTTPException(404, f"Tool '{req.tool}' not found")

    registry = get_environment_registry()
    if req.environment_id and not registry.get(req.environment_id):
        raise HTTPException(404, f"Environment '{req.environment_id}' not found")

    environment = registry.assign_tool(req.tool, req.environment_id or None)
    logger.info(
        "tool_environment_assigned",
        tool=req.tool,
        environment=req.environment_id or None,
    )
    tools_by_environment = registry._tools_by_environment()
    return {
        "ok": True,
        "tool": req.tool,
        "environment_id": req.environment_id or "",
        "environment": environment.summary(tools_by_environment.get(environment.id, [])) if environment else None,
        "tool_assignments": registry.tool_assignments,
    }
