from dataclasses import dataclass, field
from typing import Dict, List, TypedDict


class FileProgressInfo(TypedDict):
    downloaded_bytes: int
    total_bytes: int


DownloadProgressInfo = Dict[str, FileProgressInfo]


@dataclass
class ModelFileInfo:
    name: str
    downloaded: bool = False
    size: int = 0

    def to_dict(self) -> dict:
        return {
            'Name': self.name,
            'Downloaded': self.downloaded,
            'Size': self.size,
        }

    @classmethod
    def from_dict(cls, data: dict) -> 'ModelFileInfo':
        return cls(
            name=data.get('Name', ''),
            downloaded=data.get('Downloaded', False),
            size=data.get('Size', 0),
        )


@dataclass
class ModelManifest:
    name: str  # OrgName/RepoName format
    model_name: str  # Model architecture name like "qwen3-4b", "yolov12", etc.
    model_type: str  # "llm", "vlm", "embedder", etc.
    plugin_id: str = ''
    device_id: str = ''
    min_sdk_version: str = ''
    model_file: Dict[str, ModelFileInfo] = field(default_factory=dict)  # quant -> modelfile
    mmproj_file: ModelFileInfo = field(default_factory=lambda: ModelFileInfo(''))
    tokenizer_file: ModelFileInfo = field(default_factory=lambda: ModelFileInfo(''))
    extra_files: List[ModelFileInfo] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            'Name': self.name,
            'ModelName': self.model_name,
            'ModelType': self.model_type,
            'PluginId': self.plugin_id,
            'DeviceId': self.device_id,
            'MinSDKVersion': self.min_sdk_version,
            'ModelFile': {k: v.to_dict() for k, v in self.model_file.items()},
            'MMProjFile': self.mmproj_file.to_dict(),
            'TokenizerFile': self.tokenizer_file.to_dict(),
            'ExtraFiles': [f.to_dict() for f in self.extra_files],
        }

    @classmethod
    def from_dict(cls, data: dict) -> 'ModelManifest':
        return cls(
            name=data.get('Name', ''),
            model_name=data.get('ModelName', ''),
            model_type=data.get('ModelType', ''),
            plugin_id=data.get('PluginId', ''),
            device_id=data.get('DeviceId', ''),
            min_sdk_version=data.get('MinSDKVersion', ''),
            model_file={k: ModelFileInfo.from_dict(v) for k, v in data.get('ModelFile', {}).items()},
            mmproj_file=ModelFileInfo.from_dict(data.get('MMProjFile', {})),
            tokenizer_file=ModelFileInfo.from_dict(data.get('TokenizerFile', {})),
            extra_files=[ModelFileInfo.from_dict(f) for f in data.get('ExtraFiles', [])],
        )

    def get_size(self) -> int:
        total = 0
        for file_info in self.model_file.values():
            if file_info.downloaded:
                total += file_info.size
        if self.mmproj_file.downloaded:
            total += self.mmproj_file.size
        if self.tokenizer_file.downloaded:
            total += self.tokenizer_file.size
        for file_info in self.extra_files:
            if file_info.downloaded:
                total += file_info.size
        return total


@dataclass
class ModelInfo:
    repo_id: str  # Model repository ID (e.g., "microsoft/phi-3-mini-4k-instruct")
    model_name: str  # Model architecture name
    model_type: str  # Model type (e.g., "llm", "vlm", "embedder")
    size: int  # Total size in bytes
    plugin_id: str = ''  # Plugin ID
    device_id: str = ''  # Device ID
    min_sdk_version: str = ''  # Minimum SDK version required

    @classmethod
    def from_manifest(cls, manifest: ModelManifest) -> 'ModelInfo':
        return cls(
            repo_id=manifest.name,
            model_name=manifest.model_name,
            model_type=manifest.model_type,
            size=manifest.get_size(),
            plugin_id=manifest.plugin_id,
            device_id=manifest.device_id,
            min_sdk_version=manifest.min_sdk_version,
        )
