"""
LuxurAI SDK — models.py
Data classes returned by the SDK.
"""

from __future__ import annotations
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional, TYPE_CHECKING
import io

if TYPE_CHECKING:
    from .client import LuxurAI


# ─────────────────────────────────────────────
# Enums
# ─────────────────────────────────────────────

class JobStatus(str, Enum):
    PENDING   = "pending"
    QUEUED    = "queued"
    RUNNING   = "running"
    DONE      = "done"
    FAILED    = "failed"


class Resolution(str, Enum):
    """Standard output resolutions for brainAI (free-size — any value works too)."""
    SD    = "512x512"
    HD    = "1024x1024"
    WIDE  = "1024x576"
    TALL  = "576x1024"
    FHD   = "1920x1080"
    SQ_2K = "1152x1152"    # brainAI natural output size


# ─────────────────────────────────────────────
# Generated Image
# ─────────────────────────────────────────────

@dataclass
class GeneratedImage:
    """
    A generated image from LuxurAI.

    Usage:
        image = job.wait()

        # Save to file
        image.save("output.png")

        # Get raw bytes
        data = image.bytes()

        # Open with Pillow
        pil_img = image.to_pil()

        # Show in Jupyter
        image.show()
    """
    url:    str
    job_id: str
    prompt: str
    width:  int = 1024
    height: int = 1024
    _bytes: Optional[bytes] = field(default=None, repr=False)

    def bytes(self) -> bytes:
        """Download and return raw image bytes."""
        if self._bytes is None:
            import urllib.request
            with urllib.request.urlopen(self.url) as r:
                self._bytes = r.read()
        return self._bytes

    def save(self, path: str) -> None:
        """Save image to file. Format inferred from extension (.png, .jpg, .webp)."""
        with open(path, "wb") as f:
            f.write(self.bytes())
        print(f"Saved → {path}")

    def to_pil(self):
        """Return a PIL.Image object. Requires Pillow: pip install Pillow"""
        try:
            from PIL import Image
            return Image.open(io.BytesIO(self.bytes()))
        except ImportError:
            raise ImportError("Pillow required: pip install Pillow")

    def show(self) -> None:
        """Display image inline in Jupyter / IPython notebooks."""
        try:
            from IPython.display import display, Image as IPImage
            display(IPImage(data=self.bytes()))
        except ImportError:
            # Fallback: open with PIL
            self.to_pil().show()

    def __repr__(self) -> str:
        return f"GeneratedImage(job_id='{self.job_id}', size={self.width}x{self.height}, url='{self.url[:50]}...')"


# ─────────────────────────────────────────────
# Generation Job
# ─────────────────────────────────────────────

@dataclass
class GenerationJob:
    """
    A submitted generation job. Poll .status or call .wait() to get the image.

    Usage:
        job = client.generate("a sunset over mountains", wait=False)
        print(job.status)          # JobStatus.QUEUED
        print(job.queue_position)  # 2

        image = job.wait()         # blocks until done
        image.save("output.png")
    """
    id:             str
    prompt:         str
    status:         JobStatus
    cost_lc:        float
    resolution:     str
    queue_position: Optional[int]   = None
    eta_seconds:    Optional[int]   = None
    image_url:      Optional[str]   = None
    _client:        Optional[object] = field(default=None, repr=False)

    def wait(
        self,
        timeout:  int = 300,
        interval: int = 2,
        verbose:  bool = True,
    ) -> GeneratedImage:
        """
        Block until the job completes and return a GeneratedImage.

        Args:
            timeout:  Max seconds to wait (default 300)
            interval: Poll interval in seconds (default 2)
            verbose:  Print progress (default True)

        Returns:
            GeneratedImage with .save(), .bytes(), .to_pil(), .show()

        Raises:
            JobFailedError:   Job failed on server
            JobTimeoutError:  Job didn't complete within timeout
        """
        import time
        from .exceptions import JobFailedError, JobTimeoutError

        elapsed = 0
        while elapsed < timeout:
            self.refresh()

            if verbose:
                pos = f" | queue #{self.queue_position}" if self.queue_position is not None else ""
                eta = f" | eta ~{self.eta_seconds}s" if self.eta_seconds else ""
                print(f"\r  {self.status.value}{pos}{eta}    ", end="", flush=True)

            if self.status == JobStatus.DONE and self.image_url:
                if verbose:
                    print(f"\r  done in {elapsed}s                    ")
                w, h = self.resolution.split("x")
                return GeneratedImage(
                    url    = self.image_url,
                    job_id = self.id,
                    prompt = self.prompt,
                    width  = int(w),
                    height = int(h),
                )

            if self.status == JobStatus.FAILED:
                raise JobFailedError(f"Job {self.id} failed on server")

            time.sleep(interval)
            elapsed += interval

        raise JobTimeoutError(f"Job {self.id} timed out after {timeout}s")

    def refresh(self) -> "GenerationJob":
        """Fetch latest job status from API."""
        if self._client is None:
            raise RuntimeError("Job has no client — cannot refresh")
        updated = self._client.jobs.get(self.id)
        self.status         = updated.status
        self.queue_position = updated.queue_position
        self.eta_seconds    = updated.eta_seconds
        self.image_url      = updated.image_url
        return self

    def __repr__(self) -> str:
        return (
            f"GenerationJob(id='{self.id}', status={self.status.value}, "
            f"prompt='{self.prompt[:40]}...', cost={self.cost_lc}LC)"
        )


# ─────────────────────────────────────────────
# Wallet
# ─────────────────────────────────────────────

@dataclass
class WalletBalance:
    """
    LuxurCoin (LC) balance.

    Usage:
        balance = client.wallet.balance()
        print(f"{balance.lc} LC  ≈  ₹{balance.inr:.2f}")
    """
    lc:     float
    inr:    float          # lc × ₹0.25
    ledger: list = field(default_factory=list)

    def __repr__(self) -> str:
        return f"WalletBalance(lc={self.lc}, inr=₹{self.inr:.2f})"


# ─────────────────────────────────────────────
# API Key
# ─────────────────────────────────────────────

@dataclass
class APIKey:
    """
    A LuxurAI API key.

    Usage:
        keys = client.keys.list()
        for key in keys:
            print(key.prefix, key.created_at)
    """
    id:         str
    prefix:     str
    created_at: str
    is_active:  bool = True
    label:      Optional[str] = None

    def __repr__(self) -> str:
        return f"APIKey(prefix='{self.prefix}', active={self.is_active})"
