#!/usr/bin/env python
#
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from datetime import datetime
from typing import Any, Dict, Generic, List, Optional, Tuple, TypeVar, Union
from uuid import UUID, uuid4

from pydantic import BaseModel, Field, root_validator, validator
from pydantic.dataclasses import dataclass

from ._monkeypatch import _check_hotfix
from .consts import ONE_HOUR, SEVEN_DAYS
from .enums import (
    OS,
    Architecture,
    Compare,
    ContainerPermission,
    ContainerType,
    ErrorCode,
    GithubIssueSearchMatch,
    GithubIssueState,
    HeartbeatType,
    JobState,
    NodeState,
    NodeTaskState,
    PoolState,
    ScalesetState,
    StatsFormat,
    TaskDebugFlag,
    TaskFeature,
    TaskState,
    TaskType,
    VmState,
)
from .primitives import Container, PoolName, Region


class UserInfo(BaseModel):
    application_id: Optional[UUID]
    object_id: Optional[UUID]
    upn: Optional[str]


# Stores the address of a secret
class SecretAddress(BaseModel):
    # keyvault address of a secret
    url: str


T = TypeVar("T")


# This class allows us to store some data that are intended to be secret
# The secret field stores either the raw data or the address of that data
# This class allows us to maintain backward compatibility with existing
# NotificationTemplate classes
@dataclass
class SecretData(Generic[T]):
    secret: Union[T, SecretAddress]

    def __init__(self, secret: Union[T, SecretAddress]):
        try:
            self.secret = SecretAddress.parse_obj(secret)
        except Exception:
            self.secret = secret

    def __str__(self) -> str:
        return self.__repr__()

    def __repr__(self) -> str:
        if isinstance(self.secret, SecretAddress):
            return str(self.secret)
        else:
            return "[REDACTED]"


class EnumModel(BaseModel):
    @root_validator(pre=True)
    def exactly_one(cls: Any, values: Any) -> Any:
        some = []

        for field, val in values.items():
            if val is not None:
                some.append(field)

        if not some:
            raise ValueError("no variant set for enum")

        if len(some) > 1:
            raise ValueError("multiple values set for enum: %s" % some)

        return values


class Error(BaseModel):
    code: ErrorCode
    errors: List[str]


OkType = TypeVar("OkType")
Result = Union[OkType, Error]


class FileEntry(BaseModel):
    container: Container
    filename: str
    sas_url: Optional[str]


class Authentication(BaseModel):
    password: str
    public_key: str
    private_key: str


class JobConfig(BaseModel):
    project: str
    name: str
    build: str
    duration: int = Field(ge=ONE_HOUR, le=SEVEN_DAYS)


class ReproConfig(BaseModel):
    container: Container
    path: str
    duration: int = Field(ge=ONE_HOUR, le=SEVEN_DAYS)


class TaskDetails(BaseModel):
    type: TaskType
    duration: int = Field(ge=ONE_HOUR, le=SEVEN_DAYS)
    target_exe: Optional[str]
    target_env: Optional[Dict[str, str]]
    target_options: Optional[List[str]]
    target_workers: Optional[int]
    target_options_merge: Optional[bool]
    check_asan_log: Optional[bool]
    check_debugger: Optional[bool] = Field(default=True)
    check_retry_count: Optional[int] = Field(ge=0)
    check_fuzzer_help: Optional[bool]
    expect_crash_on_failure: Optional[bool]
    rename_output: Optional[bool]
    supervisor_exe: Optional[str]
    supervisor_env: Optional[Dict[str, str]]
    supervisor_options: Optional[List[str]]
    supervisor_input_marker: Optional[str]
    generator_exe: Optional[str]
    generator_env: Optional[Dict[str, str]]
    generator_options: Optional[List[str]]
    analyzer_exe: Optional[str]
    analyzer_env: Optional[Dict[str, str]]
    analyzer_options: Optional[List[str]]
    wait_for_files: Optional[ContainerType]
    stats_file: Optional[str]
    stats_format: Optional[StatsFormat]
    reboot_after_setup: Optional[bool]
    target_timeout: Optional[int] = Field(ge=1)
    ensemble_sync_delay: Optional[int]
    preserve_existing_outputs: Optional[bool]
    report_list: Optional[List[str]]
    minimized_stack_depth: Optional[int]
    coverage_filter: Optional[str]


class TaskPool(BaseModel):
    count: int
    pool_name: PoolName


class TaskVm(BaseModel):
    region: Region
    sku: str
    image: str
    count: int = Field(default=1, ge=0)
    spot_instances: bool = Field(default=False)
    reboot_after_setup: Optional[bool]


class TaskContainers(BaseModel):
    type: ContainerType
    name: Container


class TaskConfig(BaseModel):
    job_id: UUID
    prereq_tasks: Optional[List[UUID]]
    task: TaskDetails
    vm: Optional[TaskVm]
    pool: Optional[TaskPool]
    containers: List[TaskContainers]
    tags: Dict[str, str]
    debug: Optional[List[TaskDebugFlag]]
    colocate: Optional[bool]


class BlobRef(BaseModel):
    account: str
    container: Container
    name: str


class Report(BaseModel):
    input_url: Optional[str]
    input_blob: Optional[BlobRef]
    executable: str
    crash_type: str
    crash_site: str
    call_stack: List[str]
    call_stack_sha256: str
    input_sha256: str
    asan_log: Optional[str]
    task_id: UUID
    job_id: UUID
    scariness_score: Optional[int]
    scariness_description: Optional[str]
    minimized_stack: Optional[List[str]]
    minimized_stack_sha256: Optional[str]
    minimized_stack_function_names: Optional[List[str]]
    minimized_stack_function_names_sha256: Optional[str]
    minimized_stack_function_lines: Optional[List[str]]
    minimized_stack_function_lines_sha256: Optional[str]


class NoReproReport(BaseModel):
    input_sha256: str
    input_blob: Optional[BlobRef]
    executable: str
    task_id: UUID
    job_id: UUID
    tries: int
    error: Optional[str]


class CrashTestResult(BaseModel):
    crash_report: Optional[Report]
    no_repro: Optional[NoReproReport]


class RegressionReport(BaseModel):
    crash_test_result: CrashTestResult
    original_crash_test_result: Optional[CrashTestResult]


class ADODuplicateTemplate(BaseModel):
    increment: List[str]
    comment: Optional[str]
    set_state: Dict[str, str]
    ado_fields: Dict[str, str]


class ADOTemplate(BaseModel):
    base_url: str
    auth_token: SecretData[str]
    project: str
    type: str
    unique_fields: List[str]
    comment: Optional[str]
    ado_fields: Dict[str, str]
    on_duplicate: ADODuplicateTemplate

    # validator needed for backward compatibility
    @validator("auth_token", pre=True, always=True)
    def validate_auth_token(cls, v: Any) -> SecretData:
        if isinstance(v, str):
            return SecretData(secret=v)
        elif isinstance(v, SecretData):
            return v
        elif isinstance(v, dict):
            return SecretData(secret=v["secret"])
        else:
            raise TypeError(f"invalid datatype {type(v)}")


class TeamsTemplate(BaseModel):
    url: SecretData[str]

    # validator needed for backward compatibility
    @validator("url", pre=True, always=True)
    def validate_url(cls, v: Any) -> SecretData:
        if isinstance(v, str):
            return SecretData(secret=v)
        elif isinstance(v, SecretData):
            return v
        elif isinstance(v, dict):
            return SecretData(secret=v["secret"])
        else:
            raise TypeError(f"invalid datatype {type(v)}")


class ContainerDefinition(BaseModel):
    type: ContainerType
    compare: Compare
    value: int
    permissions: List[ContainerPermission]


class VmDefinition(BaseModel):
    compare: Compare
    value: int


class TaskDefinition(BaseModel):
    features: List[TaskFeature]
    containers: List[ContainerDefinition]
    monitor_queue: Optional[ContainerType]
    vm: VmDefinition


# TODO: service shouldn't pass SyncedDir, but just the url and let the agent
# come up with paths
class SyncedDir(BaseModel):
    path: str
    url: str


CONTAINER_DEF = Optional[Union[SyncedDir, List[SyncedDir]]]


class ClientCredentials(BaseModel):
    client_id: UUID
    client_secret: str


class AgentConfig(BaseModel):
    client_credentials: Optional[ClientCredentials]
    onefuzz_url: str
    pool_name: PoolName
    heartbeat_queue: Optional[str]
    instance_telemetry_key: Optional[str]
    microsoft_telemetry_key: Optional[str]
    multi_tenant_domain: Optional[str]
    instance_id: UUID


class TaskUnitConfig(BaseModel):
    instance_id: UUID
    job_id: UUID
    task_id: UUID
    task_type: TaskType
    instance_telemetry_key: Optional[str]
    microsoft_telemetry_key: Optional[str]
    heartbeat_queue: str
    # command_queue: str
    input_queue: Optional[str]
    supervisor_exe: Optional[str]
    supervisor_env: Optional[Dict[str, str]]
    supervisor_options: Optional[List[str]]
    supervisor_input_marker: Optional[str]
    target_exe: Optional[str]
    target_env: Optional[Dict[str, str]]
    target_options: Optional[List[str]]
    target_timeout: Optional[int]
    target_options_merge: Optional[bool]
    target_workers: Optional[int]
    check_asan_log: Optional[bool]
    check_debugger: Optional[bool]
    check_retry_count: Optional[int]
    check_fuzzer_help: Optional[bool]
    expect_crash_on_failure: Optional[bool]
    rename_output: Optional[bool]
    generator_exe: Optional[str]
    generator_env: Optional[Dict[str, str]]
    generator_options: Optional[List[str]]
    wait_for_files: Optional[str]
    analyzer_exe: Optional[str]
    analyzer_env: Optional[Dict[str, str]]
    analyzer_options: Optional[List[str]]
    stats_file: Optional[str]
    stats_format: Optional[StatsFormat]
    ensemble_sync_delay: Optional[int]
    report_list: Optional[List[str]]
    minimized_stack_depth: Optional[int]
    coverage_filter: Optional[str]

    # from here forwards are Container definitions.  These need to be inline
    # with TaskDefinitions and ContainerTypes
    analysis: CONTAINER_DEF
    coverage: CONTAINER_DEF
    crashes: CONTAINER_DEF
    inputs: CONTAINER_DEF
    no_repro: CONTAINER_DEF
    readonly_inputs: CONTAINER_DEF
    reports: CONTAINER_DEF
    tools: CONTAINER_DEF
    unique_inputs: CONTAINER_DEF
    unique_reports: CONTAINER_DEF
    regression_reports: CONTAINER_DEF


class Forward(BaseModel):
    src_port: int
    dst_ip: str
    dst_port: int


class ProxyConfig(BaseModel):
    url: str
    notification: str
    region: Region
    proxy_id: UUID
    forwards: List[Forward]
    instance_telemetry_key: Optional[str]
    microsoft_telemetry_key: Optional[str]
    instance_id: UUID


class ProxyHeartbeat(BaseModel):
    region: Region
    proxy_id: UUID
    forwards: List[Forward]
    timestamp: datetime = Field(default_factory=datetime.utcnow)


class Files(BaseModel):
    files: List[str]


class WorkUnit(BaseModel):
    job_id: UUID
    task_id: UUID
    task_type: TaskType

    # JSON-serialized `TaskUnitConfig`.
    config: str


class WorkSet(BaseModel):
    reboot: bool
    setup_url: str
    script: bool
    work_units: List[WorkUnit]


class WorkUnitSummary(BaseModel):
    job_id: UUID
    task_id: UUID
    task_type: TaskType


class WorkSetSummary(BaseModel):
    work_units: List[WorkUnitSummary]


class GithubIssueDuplicate(BaseModel):
    comment: Optional[str]
    labels: List[str]
    reopen: bool


class GithubIssueSearch(BaseModel):
    author: Optional[str]
    state: Optional[GithubIssueState]
    field_match: List[GithubIssueSearchMatch]
    string: str


class GithubAuth(BaseModel):
    user: str = Field(min_length=1)
    personal_access_token: str = Field(min_length=1)


class GithubIssueTemplate(BaseModel):
    auth: SecretData[GithubAuth]
    organization: str = Field(min_length=1)
    repository: str = Field(min_length=1)
    title: str = Field(min_length=1)
    body: str
    unique_search: GithubIssueSearch
    assignees: List[str]
    labels: List[str]
    on_duplicate: GithubIssueDuplicate

    # validator needed for backward compatibility
    @validator("auth", pre=True, always=True)
    def validate_auth(cls, v: Any) -> SecretData:
        def try_parse_GithubAuth(x: dict) -> Optional[GithubAuth]:
            try:
                return GithubAuth.parse_obj(x)
            except Exception:
                return None

        if isinstance(v, GithubAuth):
            return SecretData(secret=v)
        elif isinstance(v, SecretData):
            return v
        elif isinstance(v, dict):
            githubAuth = try_parse_GithubAuth(v)
            if githubAuth:
                return SecretData(secret=githubAuth)

            githubAuth = try_parse_GithubAuth(v["secret"])
            if githubAuth:
                return SecretData(secret=githubAuth)

            return SecretData(secret=v["secret"])
        else:
            raise TypeError(f"invalid datatype {type(v)}")


NotificationTemplate = Union[ADOTemplate, TeamsTemplate, GithubIssueTemplate]


class Notification(BaseModel):
    timestamp: Optional[datetime] = Field(alias="Timestamp")
    container: Container
    notification_id: UUID = Field(default_factory=uuid4)
    config: NotificationTemplate


class JobTaskInfo(BaseModel):
    task_id: UUID
    type: TaskType
    state: TaskState


class Job(BaseModel):
    timestamp: Optional[datetime] = Field(alias="Timestamp")
    job_id: UUID = Field(default_factory=uuid4)
    state: JobState = Field(default=JobState.init)
    config: JobConfig
    error: Optional[str]
    end_time: Optional[datetime] = None
    task_info: Optional[List[JobTaskInfo]]
    user_info: Optional[UserInfo]


class TaskHeartbeatEntry(BaseModel):
    task_id: UUID
    job_id: Optional[UUID]
    machine_id: UUID
    data: List[Dict[str, HeartbeatType]]


class NodeHeartbeatEntry(BaseModel):
    node_id: UUID
    data: List[Dict[str, HeartbeatType]]


class NodeCommandStopIfFree(BaseModel):
    pass


class StopNodeCommand(BaseModel):
    pass


class StopTaskNodeCommand(BaseModel):
    task_id: UUID


class NodeCommandAddSshKey(BaseModel):
    public_key: str


class NodeCommand(EnumModel):
    stop: Optional[StopNodeCommand]
    stop_task: Optional[StopTaskNodeCommand]
    add_ssh_key: Optional[NodeCommandAddSshKey]
    stop_if_free: Optional[NodeCommandStopIfFree]


class NodeCommandEnvelope(BaseModel):
    command: NodeCommand
    message_id: str


class Node(BaseModel):
    timestamp: Optional[datetime] = Field(alias="Timestamp")
    pool_name: PoolName
    pool_id: Optional[UUID]
    machine_id: UUID
    state: NodeState = Field(default=NodeState.init)
    scaleset_id: Optional[UUID] = None
    tasks: Optional[List[Tuple[UUID, NodeTaskState]]] = None
    messages: Optional[List[NodeCommand]] = None
    heartbeat: Optional[datetime]
    version: str = Field(default="1.0.0")
    reimage_requested: bool = Field(default=False)
    delete_requested: bool = Field(default=False)
    debug_keep_node: bool = Field(default=False)


class ScalesetSummary(BaseModel):
    scaleset_id: UUID
    state: ScalesetState


class NodeTasks(BaseModel):
    machine_id: UUID
    task_id: UUID
    state: NodeTaskState = Field(default=NodeTaskState.init)


class AutoScaleConfig(BaseModel):
    image: str
    max_size: int = Field(default=1000, le=1000, ge=0)  # max size of pool
    min_size: int = Field(default=0, le=1000, ge=0)  # min size of pool
    region: Optional[Region]
    scaleset_size: int  # Individual scaleset size
    spot_instances: bool = Field(default=False)
    ephemeral_os_disks: bool = Field(default=False)
    vm_sku: str

    @root_validator()
    def check_data(cls, values: Any) -> Any:
        if values["min_size"] > values["max_size"]:
            raise ValueError("The pool min_size is greater than max_size")
        return values


class Pool(BaseModel):
    timestamp: Optional[datetime] = Field(alias="Timestamp")
    name: PoolName
    pool_id: UUID = Field(default_factory=uuid4)
    os: OS
    managed: bool
    autoscale: Optional[AutoScaleConfig]
    arch: Architecture
    state: PoolState = Field(default=PoolState.init)
    client_id: Optional[UUID]
    nodes: Optional[List[Node]]
    config: Optional[AgentConfig]

    # work_queue is explicitly not saved to Tables (see save_exclude).  This is
    # intended to be used to pass the information to the CLI when the CLI asks
    # for information about what work is in the queue for the pool.
    work_queue: Optional[List[WorkSetSummary]]

    # explicitly excluded from Tables
    scaleset_summary: Optional[List[ScalesetSummary]]


class ScalesetNodeState(BaseModel):
    machine_id: UUID
    instance_id: str
    state: Optional[NodeState]


class Scaleset(BaseModel):
    timestamp: Optional[datetime] = Field(alias="Timestamp")
    pool_name: PoolName
    scaleset_id: UUID = Field(default_factory=uuid4)
    state: ScalesetState = Field(default=ScalesetState.init)
    auth: Optional[Authentication]
    vm_sku: str
    image: str
    region: Region
    size: int = Field(ge=0)
    spot_instances: bool
    ephemeral_os_disks: bool = Field(default=False)
    needs_config_update: bool = Field(default=False)
    error: Optional[Error]
    nodes: Optional[List[ScalesetNodeState]]
    client_id: Optional[UUID]
    client_object_id: Optional[UUID]
    tags: Dict[str, str] = Field(default_factory=lambda: {})


class NotificationConfig(BaseModel):
    config: NotificationTemplate


class Repro(BaseModel):
    timestamp: Optional[datetime] = Field(alias="Timestamp")
    vm_id: UUID = Field(default_factory=uuid4)
    task_id: UUID
    config: ReproConfig
    state: VmState = Field(default=VmState.init)
    auth: Optional[Authentication]
    os: OS
    error: Optional[Error]
    ip: Optional[str]
    end_time: Optional[datetime]
    user_info: Optional[UserInfo]


class ExitStatus(BaseModel):
    code: Optional[int]
    signal: Optional[int]
    success: bool


class ProcessOutput(BaseModel):
    exit_status: ExitStatus
    stderr: str
    stdout: str


class WorkerRunningEvent(BaseModel):
    task_id: UUID


class WorkerDoneEvent(BaseModel):
    task_id: UUID
    exit_status: ExitStatus
    stderr: str
    stdout: str


class WorkerEvent(EnumModel):
    done: Optional[WorkerDoneEvent]
    running: Optional[WorkerRunningEvent]


class NodeSettingUpEventData(BaseModel):
    tasks: List[UUID]


class NodeDoneEventData(BaseModel):
    error: Optional[str]
    script_output: Optional[ProcessOutput]


NodeStateData = Union[NodeSettingUpEventData, NodeDoneEventData]


class NodeStateUpdate(BaseModel):
    state: NodeState
    data: Optional[NodeStateData]

    @root_validator(pre=False, skip_on_failure=True)
    def check_data(cls, values: Any) -> Any:
        data = values.get("data")

        if data:
            state = values["state"]

            if state == NodeState.setting_up:
                if isinstance(data, NodeSettingUpEventData):
                    return values

            if state == NodeState.done:
                if isinstance(data, NodeDoneEventData):
                    return values

            raise ValueError(
                "data for node state update event does not match state = %s" % state
            )
        else:
            # For now, `data` is always optional.
            return values


class NodeEvent(EnumModel):
    state_update: Optional[NodeStateUpdate]
    worker_event: Optional[WorkerEvent]


# Temporary shim type to support hot upgrade of 1.0.0 nodes.
#
# We want future variants to use an externally-tagged repr.
NodeEventShim = Union[NodeStateUpdate, NodeEvent, WorkerEvent]


class NodeEventEnvelope(BaseModel):
    machine_id: UUID
    event: NodeEventShim


class TaskEvent(BaseModel):
    timestamp: Optional[datetime] = Field(alias="Timestamp")
    task_id: UUID
    machine_id: UUID
    event_data: WorkerEvent


class TaskEventSummary(BaseModel):
    timestamp: Optional[datetime]
    event_data: str
    event_type: str


class NodeAssignment(BaseModel):
    node_id: UUID
    scaleset_id: Optional[UUID]
    state: NodeTaskState


class Task(BaseModel):
    timestamp: Optional[datetime] = Field(alias="Timestamp")
    job_id: UUID
    task_id: UUID = Field(default_factory=uuid4)
    state: TaskState = Field(default=TaskState.init)
    os: OS
    config: TaskConfig
    error: Optional[Error]
    auth: Optional[Authentication]
    heartbeat: Optional[datetime]
    end_time: Optional[datetime]
    events: Optional[List[TaskEventSummary]]
    nodes: Optional[List[NodeAssignment]]
    user_info: Optional[UserInfo]


class InstanceConfig(BaseModel):
    # initial set of admins can only be set during deployment.
    # if admins are set, only admins can update instance configs.
    admins: Optional[List[UUID]] = None

    # if set, only admins can manage pools or scalesets
    allow_pool_management: bool = Field(default=True)

    proxy_vm_sku: str = Field(default="Standard_B2s")
    allowed_aad_tenants: List[UUID]

    def update(self, config: "InstanceConfig") -> None:
        for field in config.__fields__:
            # If no admins are set, then ignore setting admins
            if field == "admins" and self.admins is None:
                continue

            if hasattr(self, field):
                setattr(self, field, getattr(config, field))

    @validator("admins", allow_reuse=True)
    def check_admins(cls, value: Optional[List[UUID]]) -> Optional[List[UUID]]:
        if value is not None and len(value) == 0:
            raise ValueError("admins must be None or contain at least one UUID")
        return value

    # At the moment, this only checks allowed_aad_tenants, however adding
    # support for 3rd party JWT validation is anticipated in a future release.
    @root_validator()
    def check_instance_config(cls, values: Any) -> Any:
        if "allowed_aad_tenants" not in values:
            raise ValueError("missing allowed_aad_tenants")

        if not len(values["allowed_aad_tenants"]):
            raise ValueError("allowed_aad_tenants must not be empty")
        return values


_check_hotfix()
