import asyncio
from dataclasses import dataclass
from enum import Enum
import json
import logging
import os
import ssl
from typing import Any, Dict, List, Literal, Optional, TypeVar, Union

from eth_account.signers.local import LocalAccount
from pydantic import BaseModel

from tristero.api import (
    fill_order,
    fill_order_feather,
    get_quote,
    poll_updates,
    poll_updates_feather,
    get_margin_quote,
    submit_margin_order,
    submit_close_position,
    get_margin_positions,
    get_margin_position,
    poll_margin_updates,
)
from .permit2 import (
    Permit2Order,
    create_permit2_order,
    sign_margin_order,
    sign_close_position,
)
from .data import ChainID
from web3 import AsyncBaseProvider, AsyncWeb3
from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    retry_if_exception_type,
)

import certifi

logger = logging.getLogger(__name__)

P = TypeVar("P", bound=AsyncBaseProvider)


class OrderType(str, Enum):
    """Order type enum for unified API."""
    SWAP = "swap"
    FEATHER = "feather"
    MARGIN = "margin"


class WebSocketClosedError(Exception):
    pass


class SwapException(Exception):
    """Base exception for all swap-related errors."""
    pass


class StuckException(SwapException):
    """Raised when swap execution times out"""
    pass


class OrderFailedException(SwapException):
    """Order execution failed on-chain."""
    def __init__(self, message: str, order_id: str, details: dict[str, Any]):
        super().__init__(message)
        self.order_id = order_id
        self.details = details


class MarginException(Exception):
    """Exception for margin-related errors."""
    pass


class TokenSpec(BaseModel, frozen=True):
    chain_id: ChainID
    token_address: str


def make_async_w3(rpc_url: str) -> AsyncWeb3:
    insecure_ssl = os.getenv("TRISTERO_INSECURE_SSL", "").strip().lower() in {"1", "true", "yes", "y", "on"}
    if insecure_ssl:
        provider = AsyncWeb3.AsyncHTTPProvider(rpc_url, request_kwargs={"ssl": False})
        return AsyncWeb3(provider)

    ca_file = os.getenv("TRISTERO_SSL_CA_FILE", "").strip()
    if not ca_file:
        ca_file = certifi.where()

    if ca_file:
        ssl_context = ssl.create_default_context(cafile=ca_file)
        provider = AsyncWeb3.AsyncHTTPProvider(rpc_url, request_kwargs={"ssl": ssl_context})
        return AsyncWeb3(provider)

    provider = AsyncWeb3.AsyncHTTPProvider(rpc_url)
    return AsyncWeb3(provider)


@dataclass
class QuoteResult:
    """Result from get_quote operations."""
    quote_data: Dict[str, Any]
    order_type: str
    
    @property
    def is_margin(self) -> bool:
        return self.order_type == "MARGIN"
    
    @property
    def is_feather(self) -> bool:
        return self.order_type == "FEATHER"


@dataclass 
class OrderResult:
    """Result from order submission."""
    order_id: str
    order_type: str
    response: Dict[str, Any]


@dataclass
class MarginPosition:
    """Margin position data."""
    id: str
    status: str
    escrow_address: str
    filler_address: str
    taker_token_id: int
    data: Dict[str, Any]


async def wait_for_feather_completion(order_id: str):
    ws = await poll_updates_feather(order_id)
    try:
        async for msg in ws:
            if not msg:
                continue
            msg = json.loads(msg)
            status = msg['status']
            logger.info(
                {
                    "message": f"status={status}",
                    "id": "order_update",
                    "payload": msg,
                }
            )
            if status in ['Expired']:
                await ws.close()
                raise SwapException(f"Swap failed: {ws.close_reason} {msg}")
            elif status in ['Finalized']:
                await ws.close()
                return msg
        raise WebSocketClosedError("WebSocket closed without completion status")
    except Exception as e:
        if not ws.close_code:
            await ws.close()
        raise


async def wait_for_permit2_completion(order_id: str):
    ws = await poll_updates(order_id)
    try:
        async for msg in ws:
            msg = json.loads(msg)
            if isinstance(msg, dict) and ("failed" in msg or "completed" in msg):
                logger.info(
                    {
                        "message": f"failed={msg.get('failed')} completed={msg.get('completed')}",
                        "id": "order_update",
                        "payload": msg,
                    }
                )
                if msg.get("failed"):
                    await ws.close()
                    raise SwapException(f"Swap failed: {ws.close_reason} {msg}")
                elif msg.get("completed"):
                    await ws.close()
                    return msg
                continue

            if isinstance(msg, dict):
                fill_tx = msg.get("fill_tx") or msg.get("fillTx")
                if fill_tx:
                    await ws.close()
                    return msg

                if msg.get("amount_repaid") is not None or msg.get("amount_settled") is not None:
                    await ws.close()
                    return msg

            status = msg.get("status") if isinstance(msg, dict) else None
            logger.info(
                {
                    "message": f"status={status}",
                    "id": "order_update",
                    "payload": msg,
                }
            )
            if status in ["Failed", "Expired", "Rejected"]:
                await ws.close()
                raise SwapException(f"Swap failed: {ws.close_reason} {msg}")
            elif status in ["Finalized", "Filled", "Completed"]:
                await ws.close()
                return msg

        raise WebSocketClosedError("WebSocket closed without completion status")
    except Exception:
        if not ws.close_code:
            await ws.close()
        raise


async def wait_for_margin_completion(order_id: str):
    """Wait for margin order completion via WebSocket."""
    ws = await poll_margin_updates(order_id)
    try:
        async for msg in ws:
            if not msg:
                continue
            msg = json.loads(msg)
            if isinstance(msg, dict) and ("failed" in msg or "completed" in msg):
                logger.info(
                    {
                        "message": f"failed={msg.get('failed')} completed={msg.get('completed')}",
                        "id": "margin_order_update",
                        "payload": msg,
                    }
                )
                if msg.get("failed"):
                    await ws.close()
                    raise MarginException(f"Margin order failed: {msg}")
                if msg.get("completed"):
                    await ws.close()
                    return msg
                continue

            status = msg.get('status', '') if isinstance(msg, dict) else ''
            logger.info(
                {
                    "message": f"margin status={status}",
                    "id": "margin_order_update",
                    "payload": msg,
                }
            )
            if status in ['Failed', 'Expired', 'Rejected']:
                await ws.close()
                raise MarginException(f"Margin order failed: {msg}")
            elif status in ['Filled', 'Finalized', 'Completed']:
                await ws.close()
                return msg
        raise WebSocketClosedError("WebSocket closed without completion status")
    except Exception as e:
        if not ws.close_code:
            await ws.close()
        raise


async def wait_for_completion(order_id: str, order_type: Union[str, OrderType] = OrderType.SWAP):
    """
    Wait for order completion.
    
    Args:
        order_id: Order ID to wait for
        order_type: Type of order (swap, feather, margin)
    """
    if isinstance(order_type, str):
        order_type = order_type.lower()
    else:
        order_type = order_type.value
        
    if order_type == "feather":
        return await wait_for_feather_completion(order_id)
    elif order_type == "margin":
        return await wait_for_margin_completion(order_id)
    else:
        return await wait_for_permit2_completion(order_id)


@retry(
    stop=stop_after_attempt(3),
    wait=wait_exponential(multiplier=1, min=4, max=10),
    retry=retry_if_exception_type(WebSocketClosedError)
)
async def wait_for_completion_with_retry(
    order_id: str, 
    order_type: Union[str, OrderType] = OrderType.SWAP
):
    """Wait for order completion with automatic retry on WebSocket failures."""
    return await wait_for_completion(order_id, order_type)


@dataclass
class FeatherSwapResult:
    deposit_address: str
    data: Any


class FeatherException(Exception):
    pass


async def start_feather_swap(
    src_t: TokenSpec,
    dst_t: TokenSpec,
    dst_addr: str,
    raw_amount: int,
    client_id: str = ''
):
    resp = await fill_order_feather(str(src_t.chain_id.value), str(dst_t.chain_id.value), dst_addr, raw_amount, client_id)
    if resp['detail']:
        raise FeatherException(resp)
    else:
        return FeatherSwapResult(
            deposit_address=resp['deposit_address'],
            data=resp
        )


async def start_permit2_swap(
    w3: AsyncWeb3[P],
    account: LocalAccount,
    src_t: TokenSpec,
    dst_t: TokenSpec,
    raw_amount: int,
):
    order = await create_permit2_order(
        w3,
        account,
        str(src_t.chain_id.value),
        src_t.token_address,
        str(dst_t.chain_id.value),
        dst_t.token_address,
        raw_amount,
    )
    response = await fill_order(
        str(order.sig.signature.to_0x_hex()),
        order.msg.domain.model_dump(by_alias=True, mode="json"),
        order.msg.message.model_dump(by_alias=True, mode="json"),
    )
    order_id = response['id']
    return order_id


async def execute_permit2_swap(
    w3: AsyncWeb3[P],
    account: LocalAccount,
    src_t: TokenSpec,
    dst_t: TokenSpec,
    raw_amount: int,
    retry: bool = True,
    timeout: Optional[float] = None
) -> dict[str, Any]:
    """Execute and wait for swap completion."""
    order_id = await start_permit2_swap(w3, account, src_t, dst_t, raw_amount)
    logger.info(f"Swap order placed: {order_id}")

    waiter = wait_for_completion_with_retry if retry else wait_for_completion

    try:
        if timeout is None:
            return await waiter(order_id, OrderType.SWAP)

        return await asyncio.wait_for(
            waiter(order_id, OrderType.SWAP),
            timeout=timeout
        )
    except asyncio.TimeoutError as exc:
        raise StuckException(f"Swap {order_id} timed out after {timeout}s") from exc


async def execute_swap(
    w3: AsyncWeb3[P],
    account: LocalAccount,
    src_t: TokenSpec,
    dst_t: TokenSpec,
    raw_amount: int,
    retry: bool = True,
    timeout: Optional[float] = None,
) -> dict[str, Any]:
    return await execute_permit2_swap(
        w3=w3,
        account=account,
        src_t=src_t,
        dst_t=dst_t,
        raw_amount=raw_amount,
        retry=retry,
        timeout=timeout,
    )


async def request_margin_quote(
    chain_id: str,
    wallet_address: str,
    quote_currency: str,
    base_currency: str,
    leverage_ratio: int,
    collateral_amount: str,
) -> QuoteResult:
    """
    Request a quote for opening a margin position.
    
    Args:
        chain_id: Chain ID (e.g., "42161" for Arbitrum)
        wallet_address: User's wallet address
        quote_currency: Quote currency token address (e.g., USDC)
        base_currency: Base currency token address (e.g., WETH)
        leverage_ratio: Leverage ratio (e.g., 2 for 2x leverage)
        collateral_amount: Collateral amount in raw units
    
    Returns:
        QuoteResult containing quote data
    """
    quote = await get_margin_quote(
        chain_id=chain_id,
        wallet_address=wallet_address,
        quote_currency=quote_currency,
        base_currency=base_currency,
        leverage_ratio=leverage_ratio,
        collateral_amount=collateral_amount,
    )
    return QuoteResult(
        quote_data=quote,
        order_type="MARGIN",
    )


def sign_order(
    quote: QuoteResult,
    private_key: str,
) -> Dict[str, Any]:
    """
    Sign a quote to create a submittable order.
    
    Args:
        quote: Quote result from request_margin_quote
        private_key: Private key for signing
    
    Returns:
        Signed order payload
    """
    if quote.is_margin:
        return sign_margin_order(quote.quote_data, private_key)
    else:
        raise ValueError(f"Unsupported order type for signing: {quote.order_type}")


async def submit_order(signed_order: Dict[str, Any]) -> OrderResult:
    """
    Submit a signed order.
    
    Args:
        signed_order: Signed order from sign_order
    
    Returns:
        OrderResult with order_id
    """
    response = await submit_margin_order(signed_order)
    return OrderResult(
        order_id=response.get("id", response.get("order_id", "")),
        order_type="MARGIN",
        response=response,
    )


async def open_margin_position(
    private_key: str,
    chain_id: str,
    wallet_address: str,
    quote_currency: str,
    base_currency: str,
    leverage_ratio: int,
    collateral_amount: str,
    wait_for_result: bool = True,
    timeout: Optional[float] = None,
) -> Dict[str, Any]:
    """
    Open a margin position in one call (quote + sign + submit + optionally wait).
    
    Args:
        private_key: Private key for signing
        chain_id: Chain ID
        wallet_address: User's wallet address
        quote_currency: Quote currency token address
        base_currency: Base currency token address
        leverage_ratio: Leverage ratio
        collateral_amount: Collateral amount in raw units
        wait_for_result: Whether to wait for order completion
        timeout: Optional timeout in seconds
    
    Returns:
        Order result or completion result
    """
    quote = await request_margin_quote(
        chain_id=chain_id,
        wallet_address=wallet_address,
        quote_currency=quote_currency,
        base_currency=base_currency,
        leverage_ratio=leverage_ratio,
        collateral_amount=collateral_amount,
    )
    
    signed = sign_order(quote, private_key)
    result = await submit_order(signed)
    
    logger.info(f"Margin order submitted: {result.order_id}")
    
    if not wait_for_result:
        return {
            "order_id": result.order_id,
            "quote": quote.quote_data,
            "response": result.response,
        }
    
    try:
        if timeout:
            completion = await asyncio.wait_for(
                wait_for_completion_with_retry(result.order_id, OrderType.MARGIN),
                timeout=timeout
            )
        else:
            completion = await wait_for_completion_with_retry(result.order_id, OrderType.MARGIN)
        
        return {
            "order_id": result.order_id,
            "quote": quote.quote_data,
            "response": result.response,
            "completion": completion,
        }
    except asyncio.TimeoutError as exc:
        raise StuckException(f"Margin order {result.order_id} timed out after {timeout}s") from exc


async def close_margin_position(
    private_key: str,
    chain_id: str,
    position_id: int,
    escrow_contract: str,
    authorized: str,
    cash_settle: bool = False,
    fraction_bps: int = 10_000,
    deadline_seconds: int = 3600,
    wait_for_result: bool = True,
    timeout: Optional[float] = None,
) -> Dict[str, Any]:
    """
    Close a margin position.
    
    Args:
        private_key: Private key for signing
        chain_id: Chain ID
        position_id: Position ID (NFT token ID / taker_token_id)
        escrow_contract: Escrow contract address
        authorized: Authorized filler address
        cash_settle: Whether to cash settle (True) or swap settle (False)
        fraction_bps: Fraction to close in basis points (10000 = 100%)
        deadline_seconds: Deadline in seconds from now
        wait_for_result: Whether to wait for completion
        timeout: Optional timeout in seconds
    
    Returns:
        Close result
    """
    signed = sign_close_position(
        chain_id=int(chain_id),
        position_id=position_id,
        private_key=private_key,
        escrow_contract=escrow_contract,
        authorized=authorized,
        cash_settle=cash_settle,
        fraction_bps=fraction_bps,
        deadline_seconds=deadline_seconds,
    )
    
    response = await submit_close_position(signed)
    close_order = response.get("close_order")
    close_order_id = None
    if isinstance(close_order, dict):
        close_order_id = close_order.get("id")

    order_id = (
        close_order_id
        or response.get("id")
        or response.get("order_id")
        or response.get("orderId")
        or response.get("orderID")
        or ""
    )
    if not isinstance(order_id, str) or not order_id:
        raise MarginException(f"Close position response missing order id: {response}")
    
    logger.info(f"Close position order submitted: {order_id}")
    
    if not wait_for_result:
        return {
            "order_id": order_id,
            "response": response,
        }
    
    try:
        if timeout:
            completion = await asyncio.wait_for(
                wait_for_completion_with_retry(order_id, OrderType.SWAP),
                timeout=timeout
            )
        else:
            completion = await wait_for_completion_with_retry(order_id, OrderType.SWAP)
        
        return {
            "order_id": order_id,
            "response": response,
            "completion": completion,
        }
    except asyncio.TimeoutError as exc:
        raise StuckException(f"Close position {order_id} timed out after {timeout}s") from exc


async def list_margin_positions(wallet_address: str) -> List[MarginPosition]:
    """
    Get all margin positions for a wallet.
    
    Args:
        wallet_address: User's wallet address
    
    Returns:
        List of MarginPosition objects
    """
    positions = await get_margin_positions(wallet_address)
    return [
        MarginPosition(
            id=str(p.get("id", "")),
            status=p.get("status", ""),
            escrow_address=p.get("escrow_address", ""),
            filler_address=p.get("filler_address", ""),
            taker_token_id=p.get("taker_token_id", 0),
            data=p,
        )
        for p in positions
    ]


async def get_position(position_id: str) -> MarginPosition:
    """
    Get a specific margin position.
    
    Args:
        position_id: Position ID
    
    Returns:
        MarginPosition object
    """
    p = await get_margin_position(position_id)
    return MarginPosition(
        id=str(p.get("id", "")),
        status=p.get("status", ""),
        escrow_address=p.get("escrow_address", ""),
        filler_address=p.get("filler_address", ""),
        taker_token_id=p.get("taker_token_id", 0),
        data=p,
    )

