# Copyright 2026 Gaofeng Fan
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
PDK base class and loader.

PDK files are NDA-protected and NOT included in this package.
Users configure their own PDK paths via:
- Project file: ./analogpy.yaml
- User config: ~/.analogpy/config.yaml
- Environment variable: ANALOGPY_PDK_{NAME}_PATH
- Explicit path in code
"""

import os
from pathlib import Path
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field

try:
    import yaml
    HAS_YAML = True
except ImportError:
    HAS_YAML = False

from ..circuit import Net, Instance


class PDKNotFoundError(Exception):
    """Raised when a PDK cannot be found."""
    pass


@dataclass
class PDKDevice:
    """A device from a PDK with specific model."""
    name: str
    device_type: str  # nmos, pmos, etc.
    model: str
    pdk_name: str
    pins: Dict[str, Net]
    params: Dict[str, Any]

    def to_instance(self) -> Instance:
        """Convert to Instance for netlist generation."""
        return Instance(
            name=self.name,
            kind=self.device_type,
            model=self.model,
            pins=self.pins,
            params=self.params
        )


class PDK:
    """
    Process Design Kit configuration.

    Provides PDK-specific device constructors with correct model names.

    Example:
        pdk = PDK.load("tsmc28")
        mn1 = pdk.nmos("M1", d=vout, g=vin, s=gnd, b=gnd, w=1e-6, l=28e-9)
    """

    def __init__(self, name: str, config: Dict[str, Any]):
        """
        Initialize PDK from configuration.

        Args:
            name: PDK name
            config: Configuration dictionary with model_path, corners, devices
        """
        self.name = name
        self.model_path = config.get('model_path')
        self.corners = config.get('corners', {})
        self.devices = config.get('devices', {})
        self._config = config

    @classmethod
    def load(cls, name: str, path: str = None) -> 'PDK':
        """
        Load PDK by name.

        Search order:
        1. Explicit path argument
        2. Project-level analogpy.yaml
        3. User config ~/.analogpy/config.yaml
        4. Environment variable ANALOGPY_PDK_{NAME}_PATH

        Args:
            name: PDK name (e.g., "tsmc28", "gf22")
            path: Optional explicit path to PDK config or model file

        Returns:
            PDK instance

        Raises:
            PDKNotFoundError: If PDK configuration not found
        """
        # Try explicit path
        if path:
            return cls._load_from_path(name, path)

        # Try project config
        project_config = cls._find_project_config()
        if project_config and name in project_config.get('pdks', {}):
            return cls._load_from_config(name, project_config['pdks'][name])

        # Try user config
        user_config = cls._load_user_config()
        if user_config and name in user_config.get('pdks', {}):
            return cls._load_from_config(name, user_config['pdks'][name])

        # Try environment variable
        env_var = f'ANALOGPY_PDK_{name.upper()}_PATH'
        env_path = os.environ.get(env_var)
        if env_path:
            return cls._load_from_path(name, env_path)

        # Not found
        raise PDKNotFoundError(
            f"PDK '{name}' not found.\n"
            f"Configure it in one of:\n"
            f"  1. Project file: ./analogpy.yaml\n"
            f"  2. User config: ~/.analogpy/config.yaml\n"
            f"  3. Environment: {env_var}\n"
            f"  4. Explicit: PDK.load('{name}', path='/path/to/pdk')\n"
            f"\n"
            f"Example analogpy.yaml:\n"
            f"  pdks:\n"
            f"    {name}:\n"
            f"      model_path: /path/to/models.scs\n"
            f"      devices:\n"
            f"        nmos:\n"
            f"          model: nch_lvt\n"
        )

    @classmethod
    def _load_from_path(cls, name: str, path: str) -> 'PDK':
        """Load PDK from explicit path."""
        path = Path(path)

        if path.suffix in ('.yaml', '.yml') and path.exists():
            # Load from YAML config
            if not HAS_YAML:
                raise ImportError("PyYAML required for YAML config files")
            with open(path) as f:
                config = yaml.safe_load(f)
            return cls(name, config)
        else:
            # Assume it's a model file path
            config = {
                'model_path': str(path),
                'devices': {
                    'nmos': {'model': 'nch'},
                    'pmos': {'model': 'pch'},
                }
            }
            return cls(name, config)

    @classmethod
    def _load_from_config(cls, name: str, config: Dict) -> 'PDK':
        """Load PDK from configuration dictionary."""
        return cls(name, config)

    @classmethod
    def _find_project_config(cls) -> Optional[Dict]:
        """Find and load project-level config."""
        if not HAS_YAML:
            return None

        current = Path.cwd()
        for parent in [current] + list(current.parents):
            config_file = parent / "analogpy.yaml"
            if config_file.exists():
                with open(config_file) as f:
                    return yaml.safe_load(f)
        return None

    @classmethod
    def _load_user_config(cls) -> Optional[Dict]:
        """Load user-level config."""
        if not HAS_YAML:
            return None

        config_path = Path.home() / ".analogpy" / "config.yaml"
        if config_path.exists():
            with open(config_path) as f:
                return yaml.safe_load(f)
        return None

    # =========================================================================
    # Device constructors
    # =========================================================================

    def nmos(self, name: str, d: Net, g: Net, s: Net, b: Net,
             w: float, l: float, nf: int = 1, **params) -> Instance:
        """
        Create NMOS with PDK model.

        Args:
            name: Instance name
            d, g, s, b: Drain, gate, source, bulk nets
            w: Width in meters
            l: Length in meters
            nf: Number of fingers
            **params: Additional PDK-specific parameters

        Returns:
            Instance object
        """
        device_config = self.devices.get('nmos', {})
        model = device_config.get('model', 'nch')

        all_params = {"w": w, "l": l}
        if nf > 1:
            all_params["nf"] = nf
        all_params.update(params)

        return Instance(
            name=name,
            kind="nmos",
            model=model,
            pins={"d": d, "g": g, "s": s, "b": b},
            params=all_params
        )

    def pmos(self, name: str, d: Net, g: Net, s: Net, b: Net,
             w: float, l: float, nf: int = 1, **params) -> Instance:
        """
        Create PMOS with PDK model.

        Args:
            name: Instance name
            d, g, s, b: Drain, gate, source, bulk nets
            w: Width in meters
            l: Length in meters
            nf: Number of fingers
            **params: Additional PDK-specific parameters

        Returns:
            Instance object
        """
        device_config = self.devices.get('pmos', {})
        model = device_config.get('model', 'pch')

        all_params = {"w": w, "l": l}
        if nf > 1:
            all_params["nf"] = nf
        all_params.update(params)

        return Instance(
            name=name,
            kind="pmos",
            model=model,
            pins={"d": d, "g": g, "s": s, "b": b},
            params=all_params
        )

    def resistor(self, name: str, p: Net, n: Net, r: float,
                 **params) -> Instance:
        """Create resistor with PDK model."""
        device_config = self.devices.get('resistor', {})
        model = device_config.get('model', 'resistor')

        return Instance(
            name=name,
            kind="resistor",
            model=model,
            pins={"p": p, "n": n},
            params={"r": r, **params}
        )

    def capacitor(self, name: str, p: Net, n: Net, c: float,
                  **params) -> Instance:
        """Create capacitor with PDK model."""
        device_config = self.devices.get('capacitor', {})
        model = device_config.get('model', 'capacitor')

        return Instance(
            name=name,
            kind="capacitor",
            model=model,
            pins={"p": p, "n": n},
            params={"c": c, **params}
        )

    def get_corner_path(self, corner: str) -> Optional[str]:
        """Get model path for a specific corner."""
        return self.corners.get(corner)

    def __repr__(self) -> str:
        return f"PDK(name='{self.name}', model_path='{self.model_path}')"
