from pydantic import BaseModel, Field
from typing import List, Dict, Optional, Union, Any, Generic, TypeVar
from enum import Enum
from dataclasses import dataclass
import re
from openqasm3 import parser


class JobStatus(str, Enum):
    submitted = "submitted"
    ready = "ready"
    started = "started"
    canceled = "canceled"
    deleted = "deleted"
    completed = "completed"
    failed = "failed"


class NoiseModel(str, Enum):
    ideal = "ideal"
    harmony = "harmony"
    aria_1 = "aria-1"
    forte_1 = "forte-1"


class GateSet(str, Enum):
    qis = "qis"
    native = "native"


class TargetOptions(str, Enum):
    simulator = "simulator"
    harmony = "qpu.harmony"
    aria_1 = "qpu.aria-1"
    aria_2 = "qpu.aria-2"
    forte_1 = "qpu.forte-1"


class NoiseOptions(BaseModel):
    model: Optional[NoiseModel] = None
    seed: Optional[int] = None


class ErrorMitigationOptions(BaseModel):
    debiasing: Optional[bool] = None


class JobType(str, Enum):
    circuit = "ionq.circuit.v1"
    multi_circuit = "ionq.multi-circuit.v1"
    quantum_function = "quantum-function"
    optimization = "optimization"


class Ansatz(BaseModel):
    format: str = "qasm"
    data: str
    num_qubits: int

    def __init__(self, **data):

        match = re.search(r"qubit\[(\d+)\]", data["data"])
        if match is None:
            raise ValueError("Invalid QASM: No qubit declaration found")

        data["num_qubits"] = int(match.group(1))

        try:
            parser.parse(data["data"])
        except Exception as e:
            raise ValueError(f"Invalid QASM: {e}")

        super().__init__(**data)


class HamiltonianPauliTerm(BaseModel):
    pauli_string: str
    coefficient: float


@dataclass
class Hamiltonian:
    terms: List[HamiltonianPauliTerm]
    num_qubits: Optional[int]


class QuantumFunction(BaseModel):
    pass


class HamiltonianEnergyData(BaseModel):
    hamiltonian: List[HamiltonianPauliTerm]
    ansatz: Ansatz


T = TypeVar("T")


@dataclass
class WorkloadInput(Generic[T]):
    type: JobType
    input: T


class QuantumFunctionWorkload(WorkloadInput[QuantumFunction]):
    pass


class QuantumFunctionInput(BaseModel):
    data: QuantumFunction
    params: Optional[List[float]] = None


class HamiltonianEnergyQuantumFunction(QuantumFunction):
    type: str = Field(default="hamiltonian-energy", const=True)
    data: HamiltonianEnergyData


class CircuitData(BaseModel):
    qasm: Optional[str] = None
    ore: Optional[Any] = None

    def __init__(self, **data):
        super().__init__(**data)
        if self.qasm is None and self.ore is None:
            raise ValueError("Circuit must have qasm or ore")


class CircuitType(str, Enum):
    ore = "ore"
    qis = "qis"
    native = "native"


class OptimizationMethod(str, Enum):
    SPSA = "SPSA"
    Powell = "Powell"
    CG = "CG"
    BFGS = "BFGS"
    Newton_CG = "Newton-CG"
    L_BFGS_B = "L-BFGS-B"
    TNC = "TNC"
    COBYLA = "COBYLA"
    COBYQA = "COBYQA"
    SLSQP = "SLSQP"
    TRUST_CONSTR = "trust-constr"
    DOGLEG = "dogleg"
    TRUST_NCG = "trust-ncg"
    TRUST_EXACT = "trust-exact"
    TRUST_KRYLOV = "trust-krylov"


class OptimizationWorkload(BaseModel):
    quantum_function: QuantumFunction
    method: OptimizationMethod
    initial_params: Optional[List[float]] = None  # temporary setter
    log_interval: Optional[int] = 1
    options: Optional[Dict[str, Any]] = None
    maxiter: Optional[int] = None


class JobConfig(BaseModel):
    target: TargetOptions
    shots: Optional[int] = 100
    noise: Optional[NoiseOptions] = None
    error_mitigation: Optional[ErrorMitigationOptions] = None

    def __init__(self, **data):
        super().__init__(**data)
        if (
            (self.target == TargetOptions.simulator and self.noise is not None)
            or (self.target != TargetOptions.simulator)
            and self.shots is None
        ):
            raise ValueError("Shots must be set for ideal simulator jobs")


class CircuitWorkload(BaseModel):
    gateset: Optional[GateSet] = GateSet.qis
    # TODO Maybe improve typing here
    circuit: List[Dict[str, Any]]
    qubits: Optional[int] = None


class MultiCircuitWorkload(BaseModel):
    gateset: Optional[GateSet] = GateSet.qis
    # TODO Improve typing here
    circuits: List[
        Dict[str, Any]
    ]  # Union[List[CircuitList], List[CircuitOperation], None] = None
    qubits: Optional[int] = None

    # def __init__(self, **data):
    #     super().__init__(**data)
    #     if self.qubits == 0:
    #         raise ValueError("Circuit must have at least one qubit")


class JobSettings(BaseModel):
    error_mitigation: Optional[ErrorMitigationOptions] = None


class JobRequest(BaseModel, Generic[T]):
    name: Optional[str] = None
    metadata: Optional[dict] = None
    shots: Optional[int] = 100
    backend: str
    noise: Optional[NoiseOptions] = None
    type: JobType
    input: T
    settings: Optional[JobSettings] = None

    class Config:
        schema_extra = {
            "name": "My Awesome Job",
            "metadata": {"custom_key": "a string, maximum 400 chars"},
            "shots": 123,
            "backend": "simulator",
            "noise": {"model": "aria-1", "seed": 1},
            "type": "ionq.circuit.v1",
            "input": {
                "qubits": 1,
                "circuits": [{"circuit": {"gate": "h", "target": 0}}],
            },
            "settings": {
                "error_mitigation": {"debiasing": True},
            }
        }


class Progress(BaseModel):
    job_id: str
    status: JobStatus
    finished_at: Optional[str]
    value: Optional[float]
    params: Optional[List[float]]
    iteration: Optional[int]


class JobProgress(BaseModel):
    progress: List[Progress]


class Solution(BaseModel):
    minimum_value: float
    optimal_parameters: List[float]


class QuantumFunctionResult(BaseModel):
    value: Optional[float] = None
    variance: Optional[float] = None


class OptimizationResult(BaseModel):
    solution: Solution


class JobStats(BaseModel):
    gate_counts: Optional[Dict[str, int]] = {}
    qubits: Optional[int] = None
    circuits: Optional[int] = None


class JobDetails(BaseModel):
    id: str
    type: JobType
    name: Optional[str]
    status: JobStatus
    backend: Optional[TargetOptions] = TargetOptions.simulator
    noise: Optional[NoiseOptions]
    metadata: Optional[Dict[str, Union[str, int]]]
    shots: Optional[int] = 100
    settings: Optional[JobSettings]
    stats: Optional[JobStats] = JobStats()
    cost_usd: Optional[float]
    submitted_at: str
    started_at: Optional[str]
    completed_at: Optional[str]
    execution_duration_ms: Optional[int]
    predicted_execution_duration_ms: Optional[int]
    child_job_ids: Optional[List[str]]
    # TODO: type this ?
    results: Optional[Dict[str, Any]]
    failure: Optional[Dict[str, str]]
    warning: Optional[Dict[str, List[str]]]

    class Config:
        schema_extra = {
            "id": "aa54e783-0a9b-4f73-ad2f-63983b6aa4a8",
            "name": "My Awesome Job",
            "status": "completed",
            "target": "qpu.harmony",
            "noise": {"model": "harmony", "seed": 100},
            "metadata": {"custom_key": "a string, maximum 400 chars"},
            "shots": 123,
            "settings": {"error_mitigation": {"debiasing": True}},
            "stats": {
                "gate_counts": {"1q": 8, "2q": 2},
                "qubits": 4,
            },
            "submitted_at": "2025-10-29T19:24:12.280Z",
            "started_at": "2025-10-29T19:24:27.963Z",
            "completed_at": "2025-10-29T19:25:17.153Z",
            "execution_duration_ms": 13,
            "predicted_execution_duration_ms": 13,
            "child_job_ids": ["aa54e783-0a9b-4f73-ad2f-63983b6aa4a8"],
            "results": {
                "shots": {
                    "url": "/v0.4/jobs/aa54e783-0a9b-4f73-ad2f-63983b6aa4a8/results/shots"
                },
                "histogram": {
                    "url": "/v0.4/jobs/aa54e783-0a9b-4f73-ad2f-63983b6aa4a8/results/histogram"
                },
                "probabilities": {
                    "url": "/v0.4/jobs/aa54e783-0a9b-4f73-ad2f-63983b6aa4a8/results/probabilities"
                }
            },
            "failure": {"error": "An error occurred!", "code": "internal_error"},
            "warning": {
                "messages": [
                    "Warning message 1",
                    "Warning message 2",
                    "etc.",
                ]
            },
            "circuits": 1,
        }


class JobList(BaseModel):
    jobs: List[JobDetails]
    next: Optional[str]

    class Config:
        schema_extra = {
            "jobs": [
                {
                    "id": "aa54e783-0a9b-4f73-ad2f-63983b6aa4a8",
                    "name": "My Awesome Job",
                    "status": "completed",
                    "target": "qpu.harmony",
                    "noise": {"model": "harmony", "seed": 100},
                    "metadata": {"custom_key": "a string, maximum 400 chars"},
                    "shots": 123,
                    "error_mitigation": {"debias": True},
                    "gate_counts": {"1q": 8, "2q": 2},
                    "qubits": 4,
                    "cost_usd": 12.41,
                    "request": 1490932820,
                    "start": 1490932821,
                    "response": 1490932834,
                    "execution_time": 13,
                    "predicted_execution_time": 13,
                    "children": ["aa54e783-0a9b-4f73-ad2f-63983b6aa4a8"],
                    "results_url": "/v0.3/jobs/617a1f8b-59d4-435d-aa33-695433d7155e/results",
                    "failure": {
                        "error": "An error occurred!",
                        "code": "internal_error",
                    },
                    "warning": {"messages": ["<string>"]},
                    "circuits": 1,
                }
            ],
            "next": "3c90c3cc-0d44-4b50-8888-8dd25736052a",
        }


class JobUpdate(BaseModel):
    id: str
    ids: Optional[List[str]]
    status: JobStatus

    class Config:
        schema_extra = {
            "id": "aa54e783-0a9b-4f73-ad2f-63983b6aa4a8",
            "status": "ready",
        }

    def __init__(self, **data):
        super().__init__(**data)
        if self.id is None and self.ids is None:
            raise ValueError("Either 'id' or 'ids' must be set")
