"""
Figma Integration for DSS Projects

Handles Figma API communication, project/file listing, and token extraction.
Includes rate limit handling with exponential backoff.
"""

import os
import json
import asyncio
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from dataclasses import dataclass, field
import logging

logger = logging.getLogger(__name__)


# =============================================================================
# RATE LIMIT CONFIGURATION
# =============================================================================

@dataclass
class RateLimitConfig:
    """Configuration for rate limit handling."""
    max_retries: int = 5
    initial_delay: float = 1.0  # seconds
    max_delay: float = 60.0  # seconds
    backoff_factor: float = 2.0
    jitter: float = 0.1  # Random jitter factor


@dataclass
class RateLimitState:
    """Track rate limit state across requests."""
    remaining: Optional[int] = None
    reset_time: Optional[float] = None
    last_request_time: float = 0
    consecutive_429s: int = 0

    def update_from_headers(self, headers: Dict[str, str]):
        """Update state from Figma response headers."""
        if 'X-RateLimit-Remaining' in headers:
            self.remaining = int(headers['X-RateLimit-Remaining'])
        if 'X-RateLimit-Reset' in headers:
            self.reset_time = float(headers['X-RateLimit-Reset'])
        self.last_request_time = time.time()

    def get_wait_time(self) -> float:
        """Calculate wait time before next request."""
        if self.reset_time and self.remaining is not None and self.remaining <= 0:
            wait = max(0, self.reset_time - time.time())
            return wait
        return 0

    def record_429(self):
        """Record a 429 rate limit response."""
        self.consecutive_429s += 1
        self.remaining = 0

    def record_success(self):
        """Record a successful request."""
        self.consecutive_429s = 0


class FigmaRateLimitError(Exception):
    """Raised when rate limit is exceeded after retries."""
    def __init__(self, message: str, retry_after: Optional[float] = None):
        super().__init__(message)
        self.retry_after = retry_after

# Optional aiohttp import for async operations
try:
    import aiohttp
    AIOHTTP_AVAILABLE = True
except ImportError:
    AIOHTTP_AVAILABLE = False

# Fallback to requests for sync operations
try:
    import requests
    REQUESTS_AVAILABLE = True
except ImportError:
    REQUESTS_AVAILABLE = False


@dataclass
class FigmaAPIConfig:
    """Figma API configuration."""
    token: str
    base_url: str = "https://api.figma.com/v1"
    timeout: int = 30
    rate_limit: RateLimitConfig = field(default_factory=RateLimitConfig)


@dataclass
class FigmaStyleData:
    """Extracted style data from Figma."""
    colors: Dict[str, Any] = field(default_factory=dict)
    typography: Dict[str, Any] = field(default_factory=dict)
    effects: Dict[str, Any] = field(default_factory=dict)
    grids: Dict[str, Any] = field(default_factory=dict)
    variables: Dict[str, Any] = field(default_factory=dict)
    raw_styles: Dict[str, Any] = field(default_factory=dict)


class FigmaProjectSync:
    """
    Synchronize design tokens from Figma projects/files.

    Supports:
    - Listing project files
    - Extracting styles from files
    - Converting to DSS token format
    """

    def __init__(self, token: Optional[str] = None, rate_limit_config: Optional[RateLimitConfig] = None):
        """
        Initialize Figma sync.

        Args:
            token: Figma personal access token. Falls back to FIGMA_TOKEN env var.
            rate_limit_config: Optional rate limit configuration.
        """
        self.token = token or os.environ.get("FIGMA_TOKEN", "")
        if not self.token:
            raise ValueError("Figma token required. Set FIGMA_TOKEN env var or pass token parameter.")

        self.config = FigmaAPIConfig(
            token=self.token,
            rate_limit=rate_limit_config or RateLimitConfig()
        )
        self._session: Optional[aiohttp.ClientSession] = None
        self._rate_limit_state = RateLimitState()

    @property
    def headers(self) -> Dict[str, str]:
        """API request headers."""
        return {"X-Figma-Token": self.token}

    # =========================================================================
    # Rate Limit Handling
    # =========================================================================

    def _calculate_backoff_delay(self, attempt: int, retry_after: Optional[float] = None) -> float:
        """Calculate delay with exponential backoff and jitter."""
        import random

        config = self.config.rate_limit

        # Use Retry-After header if available
        if retry_after:
            base_delay = retry_after
        else:
            base_delay = config.initial_delay * (config.backoff_factor ** attempt)

        # Cap at max delay
        delay = min(base_delay, config.max_delay)

        # Add jitter
        jitter = delay * config.jitter * random.random()
        return delay + jitter

    def _request_with_retry(
        self,
        method: str,
        url: str,
        **kwargs
    ) -> requests.Response:
        """
        Make HTTP request with rate limit retry logic.

        Args:
            method: HTTP method (get, post, etc.)
            url: Request URL
            **kwargs: Additional request arguments

        Returns:
            Response object

        Raises:
            FigmaRateLimitError: If rate limit exceeded after all retries
            requests.HTTPError: For other HTTP errors
        """
        if not REQUESTS_AVAILABLE:
            raise ImportError("requests library required for sync operations")

        config = self.config.rate_limit
        last_error = None

        # Pre-emptive wait if we know rate limit is exhausted
        wait_time = self._rate_limit_state.get_wait_time()
        if wait_time > 0:
            logger.info(f"Rate limit: waiting {wait_time:.1f}s before request")
            time.sleep(wait_time)

        for attempt in range(config.max_retries + 1):
            try:
                # Make request
                response = requests.request(
                    method,
                    url,
                    headers=self.headers,
                    timeout=self.config.timeout,
                    **kwargs
                )

                # Update rate limit state from headers
                self._rate_limit_state.update_from_headers(dict(response.headers))

                # Handle rate limit (429)
                if response.status_code == 429:
                    self._rate_limit_state.record_429()

                    # Get retry-after from header
                    retry_after = None
                    if 'Retry-After' in response.headers:
                        try:
                            retry_after = float(response.headers['Retry-After'])
                        except ValueError:
                            pass

                    if attempt < config.max_retries:
                        delay = self._calculate_backoff_delay(attempt, retry_after)
                        logger.warning(
                            f"Rate limited (429). Attempt {attempt + 1}/{config.max_retries + 1}. "
                            f"Waiting {delay:.1f}s before retry..."
                        )
                        time.sleep(delay)
                        continue
                    else:
                        raise FigmaRateLimitError(
                            f"Rate limit exceeded after {config.max_retries} retries",
                            retry_after=retry_after
                        )

                # Success
                self._rate_limit_state.record_success()
                response.raise_for_status()
                return response

            except requests.exceptions.RequestException as e:
                last_error = e
                if attempt < config.max_retries:
                    delay = self._calculate_backoff_delay(attempt)
                    logger.warning(
                        f"Request failed: {e}. Attempt {attempt + 1}/{config.max_retries + 1}. "
                        f"Waiting {delay:.1f}s before retry..."
                    )
                    time.sleep(delay)
                    continue
                raise

        # Should not reach here, but just in case
        if last_error:
            raise last_error
        raise RuntimeError("Unexpected state in retry loop")

    async def _request_with_retry_async(
        self,
        method: str,
        url: str,
        **kwargs
    ) -> Tuple[int, Dict[str, Any]]:
        """
        Make async HTTP request with rate limit retry logic.

        Returns:
            Tuple of (status_code, response_json)
        """
        if not AIOHTTP_AVAILABLE:
            raise ImportError("aiohttp library required for async operations")

        import random
        config = self.config.rate_limit
        session = await self._get_session()
        last_error = None

        # Pre-emptive wait if we know rate limit is exhausted
        wait_time = self._rate_limit_state.get_wait_time()
        if wait_time > 0:
            logger.info(f"Rate limit: waiting {wait_time:.1f}s before request")
            await asyncio.sleep(wait_time)

        for attempt in range(config.max_retries + 1):
            try:
                async with session.request(method, url, **kwargs) as response:
                    # Update rate limit state from headers
                    self._rate_limit_state.update_from_headers(dict(response.headers))

                    # Handle rate limit (429)
                    if response.status == 429:
                        self._rate_limit_state.record_429()

                        retry_after = None
                        if 'Retry-After' in response.headers:
                            try:
                                retry_after = float(response.headers['Retry-After'])
                            except ValueError:
                                pass

                        if attempt < config.max_retries:
                            delay = self._calculate_backoff_delay(attempt, retry_after)
                            logger.warning(
                                f"Rate limited (429). Attempt {attempt + 1}/{config.max_retries + 1}. "
                                f"Waiting {delay:.1f}s before retry..."
                            )
                            await asyncio.sleep(delay)
                            continue
                        else:
                            raise FigmaRateLimitError(
                                f"Rate limit exceeded after {config.max_retries} retries",
                                retry_after=retry_after
                            )

                    # Success
                    self._rate_limit_state.record_success()
                    data = await response.json()
                    return response.status, data

            except aiohttp.ClientError as e:
                last_error = e
                if attempt < config.max_retries:
                    delay = self._calculate_backoff_delay(attempt)
                    logger.warning(
                        f"Request failed: {e}. Attempt {attempt + 1}/{config.max_retries + 1}. "
                        f"Waiting {delay:.1f}s before retry..."
                    )
                    await asyncio.sleep(delay)
                    continue
                raise

        if last_error:
            raise last_error
        raise RuntimeError("Unexpected state in retry loop")

    def get_rate_limit_status(self) -> Dict[str, Any]:
        """Get current rate limit status."""
        state = self._rate_limit_state
        return {
            "remaining": state.remaining,
            "reset_time": state.reset_time,
            "reset_in_seconds": max(0, state.reset_time - time.time()) if state.reset_time else None,
            "consecutive_429s": state.consecutive_429s,
            "last_request_time": state.last_request_time,
        }

    # =========================================================================
    # Sync API (uses requests)
    # =========================================================================

    def list_project_files(self, project_id: str) -> Dict[str, Any]:
        """
        List all files in a Figma project (sync).

        Args:
            project_id: Figma project ID

        Returns:
            Dict with project name and files list
        """
        url = f"{self.config.base_url}/projects/{project_id}/files"
        response = self._request_with_retry("GET", url)
        data = response.json()

        return {
            "project_name": data.get("name", ""),
            "files": [
                {
                    "key": f.get("key"),
                    "name": f.get("name"),
                    "thumbnail_url": f.get("thumbnail_url"),
                    "last_modified": f.get("last_modified"),
                }
                for f in data.get("files", [])
            ]
        }

    def list_team_projects(self, team_id: str) -> Dict[str, Any]:
        """
        List all projects in a Figma team (sync).

        Args:
            team_id: Figma team ID

        Returns:
            Dict with team projects
        """
        url = f"{self.config.base_url}/teams/{team_id}/projects"
        response = self._request_with_retry("GET", url)
        data = response.json()

        return {
            "team_name": data.get("name", ""),
            "projects": [
                {
                    "id": p.get("id"),
                    "name": p.get("name"),
                }
                for p in data.get("projects", [])
            ]
        }

    def discover_team_structure(self, team_id: str) -> Dict[str, Any]:
        """
        Discover the full structure of a Figma team.

        Returns team projects and their files, identifying the UIKit reference file.
        Uses rate limit handling for all API calls.

        Args:
            team_id: Figma team ID

        Returns:
            Dict with full team structure including identified uikit file
        """

        # Get all projects in team
        team_data = self.list_team_projects(team_id)

        result = {
            "team_id": team_id,
            "team_name": team_data.get("team_name", ""),
            "projects": [],
            "uikit": None,  # Will be populated if found
        }

        # For each project, get files
        for project in team_data.get("projects", []):
            project_id = project["id"]
            project_name = project["name"]

            try:
                project_files = self.list_project_files(project_id)

                project_data = {
                    "id": project_id,
                    "name": project_name,
                    "files": project_files.get("files", []),
                }
                result["projects"].append(project_data)

                # Search for UIKit file in this project
                for file in project_data["files"]:
                    file_name_lower = file.get("name", "").lower()
                    # Look for common UIKit naming patterns
                    if any(pattern in file_name_lower for pattern in [
                        "uikit", "ui-kit", "ui kit",
                        "design system", "design-system",
                        "tokens", "foundations",
                        "core", "base"
                    ]):
                        # Prefer exact "uikit" match
                        is_better_match = (
                            result["uikit"] is None or
                            "uikit" in file_name_lower and "uikit" not in result["uikit"]["name"].lower()
                        )
                        if is_better_match:
                            result["uikit"] = {
                                "key": file["key"],
                                "name": file["name"],
                                "project_id": project_id,
                                "project_name": project_name,
                            }

            except Exception as e:
                logger.warning(f"Failed to get files for project {project_name}: {e}")

        return result

    def find_uikit_file(self, team_id: str) -> Optional[Dict[str, Any]]:
        """
        Find the UIKit reference file in a team.

        Searches all projects for a file named 'uikit' or similar.

        Args:
            team_id: Figma team ID

        Returns:
            Dict with uikit file info or None if not found
        """
        structure = self.discover_team_structure(team_id)
        return structure.get("uikit")

    def get_file_styles(self, file_key: str) -> FigmaStyleData:
        """
        Extract styles from a Figma file (sync).

        Uses rate limit handling with exponential backoff for all API calls.

        Args:
            file_key: Figma file key

        Returns:
            FigmaStyleData with extracted styles
        """
        # Get file data with retry
        url = f"{self.config.base_url}/files/{file_key}"
        response = self._request_with_retry("GET", url)
        file_data = response.json()

        # Get styles with retry
        styles_url = f"{self.config.base_url}/files/{file_key}/styles"
        styles_response = self._request_with_retry("GET", styles_url)
        styles_data = styles_response.json()

        # Get variables (if available - newer Figma API)
        variables = {}
        try:
            vars_url = f"{self.config.base_url}/files/{file_key}/variables/local"
            vars_response = self._request_with_retry("GET", vars_url)
            variables = vars_response.json()
        except FigmaRateLimitError:
            # Re-raise rate limit errors
            raise
        except Exception as e:
            logger.debug(f"Variables not available for file {file_key}: {e}")

        return self._parse_styles(file_data, styles_data, variables)

    # =========================================================================
    # Async API (uses aiohttp)
    # =========================================================================

    async def _get_session(self) -> aiohttp.ClientSession:
        """Get or create aiohttp session."""
        if not AIOHTTP_AVAILABLE:
            raise ImportError("aiohttp library required for async operations")

        if self._session is None or self._session.closed:
            timeout = aiohttp.ClientTimeout(total=self.config.timeout)
            self._session = aiohttp.ClientSession(
                headers=self.headers,
                timeout=timeout
            )
        return self._session

    async def close(self):
        """Close the aiohttp session."""
        if self._session and not self._session.closed:
            await self._session.close()

    async def list_project_files_async(self, project_id: str) -> Dict[str, Any]:
        """List all files in a Figma project (async) with rate limit handling."""
        url = f"{self.config.base_url}/projects/{project_id}/files"
        status, data = await self._request_with_retry_async("GET", url)

        if status != 200:
            raise ValueError(f"Failed to list project files: status {status}")

        return {
            "project_name": data.get("name", ""),
            "files": [
                {
                    "key": f.get("key"),
                    "name": f.get("name"),
                    "thumbnail_url": f.get("thumbnail_url"),
                    "last_modified": f.get("last_modified"),
                }
                for f in data.get("files", [])
            ]
        }

    async def get_file_styles_async(self, file_key: str) -> FigmaStyleData:
        """Extract styles from a Figma file (async) with rate limit handling.

        Note: Requests are made sequentially to respect rate limits.
        """
        # Get file data
        file_url = f"{self.config.base_url}/files/{file_key}"
        file_status, file_data = await self._request_with_retry_async("GET", file_url)

        if file_status != 200:
            raise ValueError(f"Failed to fetch file {file_key}: status {file_status}")

        # Get styles
        styles_url = f"{self.config.base_url}/files/{file_key}/styles"
        styles_status, styles_data = await self._request_with_retry_async("GET", styles_url)

        if styles_status != 200:
            styles_data = {}

        # Get variables (if available - newer Figma API)
        variables = {}
        try:
            vars_url = f"{self.config.base_url}/files/{file_key}/variables/local"
            vars_status, vars_data = await self._request_with_retry_async("GET", vars_url)
            if vars_status == 200:
                variables = vars_data
        except FigmaRateLimitError:
            raise
        except Exception as e:
            logger.debug(f"Variables not available for file {file_key}: {e}")

        return self._parse_styles(file_data, styles_data, variables)

    async def sync_project_files_async(
        self,
        project_id: str,
        file_keys: Optional[List[str]] = None
    ) -> Dict[str, FigmaStyleData]:
        """
        Sync styles from multiple files in a project (async).

        Args:
            project_id: Figma project ID
            file_keys: Optional list of specific file keys. If None, syncs all.

        Returns:
            Dict mapping file keys to their extracted styles
        """
        # Get project files if not specified
        if file_keys is None:
            project_data = await self.list_project_files_async(project_id)
            file_keys = [f["key"] for f in project_data["files"]]

        # Fetch styles from all files in parallel
        tasks = [self.get_file_styles_async(key) for key in file_keys]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        styles_map = {}
        for key, result in zip(file_keys, results):
            if isinstance(result, Exception):
                logger.error(f"Failed to sync file {key}: {result}")
            else:
                styles_map[key] = result

        return styles_map

    # =========================================================================
    # Style Parsing
    # =========================================================================

    def _parse_styles(
        self,
        file_data: Dict[str, Any],
        styles_data: Dict[str, Any],
        variables: Dict[str, Any]
    ) -> FigmaStyleData:
        """Parse Figma API responses into FigmaStyleData."""
        result = FigmaStyleData()

        # Parse document styles
        document = file_data.get("document", {})
        global_styles = file_data.get("styles", {})

        # Extract colors from styles
        result.colors = self._extract_colors(global_styles, document)

        # Extract typography
        result.typography = self._extract_typography(global_styles, document)

        # Extract effects (shadows, blurs)
        result.effects = self._extract_effects(global_styles, document)

        # Extract variables (new Figma variables API)
        if variables:
            result.variables = self._extract_variables(variables)

        # Store raw styles for reference
        result.raw_styles = {
            "global_styles": global_styles,
            "meta": styles_data.get("meta", {}),
        }

        return result

    def _extract_colors(
        self,
        global_styles: Dict[str, Any],
        document: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Extract color styles."""
        colors = {}

        for style_id, style in global_styles.items():
            if style.get("styleType") == "FILL":
                name = style.get("name", style_id)
                # Normalize name to token path format
                token_name = self._normalize_name(name)
                colors[token_name] = {
                    "figma_id": style_id,
                    "name": name,
                    "description": style.get("description", ""),
                }

        return colors

    def _extract_typography(
        self,
        global_styles: Dict[str, Any],
        document: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Extract typography styles."""
        typography = {}

        for style_id, style in global_styles.items():
            if style.get("styleType") == "TEXT":
                name = style.get("name", style_id)
                token_name = self._normalize_name(name)
                typography[token_name] = {
                    "figma_id": style_id,
                    "name": name,
                    "description": style.get("description", ""),
                }

        return typography

    def _extract_effects(
        self,
        global_styles: Dict[str, Any],
        document: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Extract effect styles (shadows, blurs)."""
        effects = {}

        for style_id, style in global_styles.items():
            if style.get("styleType") == "EFFECT":
                name = style.get("name", style_id)
                token_name = self._normalize_name(name)
                effects[token_name] = {
                    "figma_id": style_id,
                    "name": name,
                    "description": style.get("description", ""),
                }

        return effects

    def _extract_variables(self, variables_data: Dict[str, Any]) -> Dict[str, Any]:
        """Extract Figma variables (new API)."""
        variables = {}

        meta = variables_data.get("meta", {})
        var_collections = meta.get("variableCollections", {})
        var_values = meta.get("variables", {})

        for var_id, var_data in var_values.items():
            name = var_data.get("name", var_id)
            resolved_type = var_data.get("resolvedType", "")

            token_name = self._normalize_name(name)
            variables[token_name] = {
                "figma_id": var_id,
                "name": name,
                "type": resolved_type,
                "description": var_data.get("description", ""),
                "values": var_data.get("valuesByMode", {}),
            }

        return variables

    def _normalize_name(self, name: str) -> str:
        """Normalize Figma style name to token path format."""
        # Convert "Colors/Primary/500" -> "colors.primary.500"
        # Convert "Typography/Heading/H1" -> "typography.heading.h1"
        normalized = name.lower()
        normalized = normalized.replace("/", ".")
        normalized = normalized.replace(" ", "-")
        normalized = normalized.replace("--", "-")
        return normalized

    # =========================================================================
    # Token Conversion
    # =========================================================================

    def to_dss_tokens(self, style_data: FigmaStyleData) -> Dict[str, Any]:
        """
        Convert FigmaStyleData to DSS token format.

        Returns a dict compatible with DSS TokenCollection.
        """
        tokens = {
            "source": "figma",
            "timestamp": datetime.now().isoformat(),
            "tokens": {}
        }

        # Add color tokens
        for path, data in style_data.colors.items():
            tokens["tokens"][f"color.{path}"] = {
                "value": None,  # Will be resolved from Figma node data
                "type": "color",
                "source": "figma",
                "metadata": data,
            }

        # Add typography tokens
        for path, data in style_data.typography.items():
            tokens["tokens"][f"typography.{path}"] = {
                "value": None,
                "type": "typography",
                "source": "figma",
                "metadata": data,
            }

        # Add effect tokens
        for path, data in style_data.effects.items():
            tokens["tokens"][f"effect.{path}"] = {
                "value": None,
                "type": "effect",
                "source": "figma",
                "metadata": data,
            }

        # Add variables (these have actual values)
        for path, data in style_data.variables.items():
            var_type = data.get("type", "").lower()
            if var_type == "color":
                prefix = "color"
            elif var_type == "float":
                prefix = "size"
            elif var_type == "string":
                prefix = "string"
            else:
                prefix = "var"

            tokens["tokens"][f"{prefix}.{path}"] = {
                "value": data.get("values", {}),
                "type": var_type or "unknown",
                "source": "figma-variable",
                "metadata": data,
            }

        return tokens

    def save_tokens(
        self,
        style_data: FigmaStyleData,
        output_path: Path,
        format: str = "json"
    ) -> Path:
        """
        Save extracted tokens to file.

        Args:
            style_data: Extracted Figma styles
            output_path: Directory to save to
            format: Output format (json, raw)

        Returns:
            Path to saved file
        """
        output_path = Path(output_path)
        output_path.mkdir(parents=True, exist_ok=True)

        if format == "json":
            tokens = self.to_dss_tokens(style_data)
            file_path = output_path / "figma-tokens.json"
            with open(file_path, "w") as f:
                json.dump(tokens, f, indent=2)
        elif format == "raw":
            file_path = output_path / "figma-raw.json"
            with open(file_path, "w") as f:
                json.dump({
                    "colors": style_data.colors,
                    "typography": style_data.typography,
                    "effects": style_data.effects,
                    "variables": style_data.variables,
                }, f, indent=2)
        else:
            raise ValueError(f"Unknown format: {format}")

        return file_path
