from __future__ import annotations

import base64
from dataclasses import dataclass
from pathlib import Path
from time import perf_counter
from typing import Any
from uuid import uuid4

from cryptography.fernet import Fernet
from django.conf import settings
from django.core.management.base import BaseCommand, CommandError
from django.db import transaction
from django.test.utils import override_settings

from apps.ai.aliases import resolve_alias
from apps.ai.models import AiModel, ModelAlias
from apps.ai.providers import get_provider
from apps.ai.providers.openai_compatible import ChatCompletionsProvider


DEFAULT_TITLE_PROMPT = "为一件简约白色T恤生成3个中文电商标题"
RECORDED_TITLE_CONTENT = (
    "1. 简约白T清爽百搭日常通勤短袖\n"
    "2. 纯净白色基础款T恤舒适亲肤上衣\n"
    "3. 夏季百搭白T休闲宽松男女同款"
)
RECORDED_IMAGE_BYTES = b"recorded-image-bytes"


@dataclass(frozen=True)
class SmokeResult:
    mode: str
    operation: str
    alias: str
    model_used: str
    elapsed_ms: int
    title_count: int | None = None
    first_title: str = ""
    image_bytes: int | None = None


class RecordedResponse:
    def __init__(self, payload: dict[str, Any]) -> None:
        self.payload = payload

    def json(self) -> dict[str, Any]:
        return self.payload

    def raise_for_status(self) -> None:
        return None


class RecordedSession:
    def __init__(self, payload: dict[str, Any]) -> None:
        self.payload = payload
        self.posts: list[dict[str, Any]] = []
        self.trust_env = True

    def post(self, url: str, **kwargs: Any) -> RecordedResponse:
        sanitized_kwargs = {
            key: value for key, value in kwargs.items() if key != "headers"
        }
        self.posts.append({"url": url, **sanitized_kwargs})
        return RecordedResponse(self.payload)


class Command(BaseCommand):
    help = "Run a minimal AI generation smoke through alias resolution and provider code."

    def add_arguments(self, parser):
        parser.add_argument("operation", choices=("title", "image"))
        parser.add_argument("--alias", help="Alias to resolve. Defaults to operation default.")
        parser.add_argument("--prompt", default=DEFAULT_TITLE_PROMPT)
        parser.add_argument("--resolution", default="1K")
        parser.add_argument("--image-file", help="Input image path for image/edit smoke.")
        parser.add_argument("--image-mime-type", default="image/png")
        parser.add_argument(
            "--recorded",
            action="store_true",
            help="Use a recorded provider response and temporary rolled-back config.",
        )

    def handle(self, *args, **options):
        if options["operation"] == "title":
            result = (
                self._run_recorded_title(options)
                if options["recorded"]
                else self._run_real_title(options)
            )
        else:
            result = (
                self._run_recorded_image(options)
                if options["recorded"]
                else self._run_real_image(options)
            )
        self._write_result(result)

    def _run_real_title(self, options) -> SmokeResult:
        if not settings.AI_KEY_ENCRYPTION_KEY:
            raise CommandError("AI_KEY_ENCRYPTION_KEY is not configured.")
        return run_title_smoke(
            mode="real",
            alias=options["alias"],
            prompt=options["prompt"],
            resolution=options["resolution"],
            provider=None,
        )

    def _run_recorded_title(self, options) -> SmokeResult:
        alias = options["alias"] or f"t-104-recorded-title-{uuid4().hex[:8]}"
        encryption_key = Fernet.generate_key().decode("ascii")
        with override_settings(AI_KEY_ENCRYPTION_KEY=encryption_key):
            with transaction.atomic():
                ai_model = AiModel(
                    name=f"T-104 recorded title {uuid4().hex[:8]}",
                    url="https://recorded.invalid/v1",
                    model="recorded-title-model",
                    api_type=AiModel.ApiType.CHAT,
                    capabilities=["text"],
                    timeout_seconds=0,
                    connect_timeout_seconds=30,
                )
                ai_model.set_api_key("sk-recorded-placeholder")
                ai_model.save()
                ModelAlias.objects.create(
                    operation_type=ModelAlias.OperationType.TITLE,
                    alias=alias,
                    ai_model=ai_model,
                )
                result = run_title_smoke(
                    mode="recorded",
                    alias=alias,
                    prompt=options["prompt"],
                    resolution=options["resolution"],
                    provider=ChatCompletionsProvider(
                        session=RecordedSession(recorded_title_payload())
                    ),
                )
                transaction.set_rollback(True)
                return result

    def _run_real_image(self, options) -> SmokeResult:
        if not settings.AI_KEY_ENCRYPTION_KEY:
            raise CommandError("AI_KEY_ENCRYPTION_KEY is not configured.")
        return run_image_smoke(
            mode="real",
            alias=options["alias"],
            prompt=options["prompt"],
            resolution=options["resolution"],
            image=read_image_file(options["image_file"]),
            image_mime_type=options["image_mime_type"],
            provider=None,
        )

    def _run_recorded_image(self, options) -> SmokeResult:
        alias = options["alias"] or f"t-105-recorded-image-{uuid4().hex[:8]}"
        encryption_key = Fernet.generate_key().decode("ascii")
        with override_settings(AI_KEY_ENCRYPTION_KEY=encryption_key):
            with transaction.atomic():
                ai_model = AiModel(
                    name=f"T-105 recorded image {uuid4().hex[:8]}",
                    url="https://recorded.invalid/v1/chat/completions",
                    model="recorded-image-model",
                    api_type=AiModel.ApiType.CHAT,
                    capabilities=["image", "vision"],
                    timeout_seconds=0,
                    connect_timeout_seconds=30,
                )
                ai_model.set_api_key("sk-recorded-placeholder")
                ai_model.save()
                ModelAlias.objects.create(
                    operation_type=ModelAlias.OperationType.IMAGE,
                    alias=alias,
                    ai_model=ai_model,
                )
                result = run_image_smoke(
                    mode="recorded",
                    alias=alias,
                    prompt=options["prompt"],
                    resolution=options["resolution"],
                    image=b"recorded-input-image",
                    image_mime_type=options["image_mime_type"],
                    provider=ChatCompletionsProvider(
                        session=RecordedSession(recorded_image_payload())
                    ),
                )
                transaction.set_rollback(True)
                return result

    def _write_result(self, result: SmokeResult) -> None:
        self.stdout.write(self.style.SUCCESS("Smoke generation OK"))
        self.stdout.write(f"mode={result.mode}")
        self.stdout.write(f"operation={result.operation}")
        self.stdout.write(f"alias={result.alias}")
        self.stdout.write(f"model_used={result.model_used}")
        self.stdout.write(f"elapsed_ms={result.elapsed_ms}")
        if result.title_count is not None:
            self.stdout.write(f"title_count={result.title_count}")
            self.stdout.write(f"first_title={result.first_title}")
        if result.image_bytes is not None:
            self.stdout.write(f"image_bytes={result.image_bytes}")


def run_title_smoke(
    *,
    mode: str,
    alias: str | None,
    prompt: str,
    resolution: str,
    provider,
) -> SmokeResult:
    started = perf_counter()
    model = resolve_alias(ModelAlias.OperationType.TITLE, alias)
    selected_provider = provider or get_provider(model.api_type, model.url)
    generation = selected_provider.generate_text(
        prompt,
        model,
        resolution=resolution,
        parameters={"temperature": 0},
    )
    elapsed_ms = int((perf_counter() - started) * 1000)
    return SmokeResult(
        mode=mode,
        operation=ModelAlias.OperationType.TITLE,
        alias=alias or "<default:title>",
        model_used=generation.model_used,
        elapsed_ms=elapsed_ms,
        title_count=len(generation.titles),
        first_title=generation.titles[0] if generation.titles else generation.text,
    )


def run_image_smoke(
    *,
    mode: str,
    alias: str | None,
    prompt: str,
    resolution: str,
    image: bytes | None,
    image_mime_type: str,
    provider,
) -> SmokeResult:
    started = perf_counter()
    model = resolve_alias(ModelAlias.OperationType.IMAGE, alias)
    selected_provider = provider or get_provider(model.api_type, model.url)
    generation = selected_provider.generate_image(
        prompt,
        model,
        image=image,
        image_mime_type=image_mime_type,
        resolution=resolution,
        parameters={"temperature": 0},
    )
    elapsed_ms = int((perf_counter() - started) * 1000)
    return SmokeResult(
        mode=mode,
        operation=ModelAlias.OperationType.IMAGE,
        alias=alias or "<default:image>",
        model_used=generation.model_used,
        elapsed_ms=elapsed_ms,
        image_bytes=len(generation.image),
    )


def recorded_title_payload() -> dict[str, Any]:
    return {"choices": [{"message": {"content": RECORDED_TITLE_CONTENT}}]}


def recorded_image_payload() -> dict[str, Any]:
    encoded = base64.b64encode(RECORDED_IMAGE_BYTES).decode("ascii")
    return {
        "choices": [
            {
                "message": {
                    "content": [
                        {"type": "text", "text": "done"},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{encoded}",
                            },
                        },
                    ]
                }
            }
        ]
    }


def read_image_file(path: str | None) -> bytes | None:
    if not path:
        return None
    try:
        return Path(path).read_bytes()
    except OSError as exc:
        raise CommandError(f"Unable to read image file: {path}") from exc