import asyncio
import ipaddress
import json
import socket
from urllib.parse import urlparse

import httpx
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from typing import Any, Dict

from app.api.v1.auth import get_current_user
from app.api.v1.user_config import _load_ai_configs, _get_default_providers

router = APIRouter()

class ProxyRequest(BaseModel):
    url: str
    headers: Dict[str, str] = {}
    body: Dict[str, Any] = {}


# Ranges that are never legitimate AI endpoints: cloud metadata + this host.
_BLOCKED_NETWORKS = [
    ipaddress.ip_network("169.254.0.0/16"),   # link-local / cloud metadata
    ipaddress.ip_network("0.0.0.0/8"),
]
# Private ranges: only reachable when the target host is one the user has
# explicitly configured as an AI provider (e.g. local Ollama/LM Studio).
_PRIVATE_NETWORKS = [
    ipaddress.ip_network("10.0.0.0/8"),
    ipaddress.ip_network("172.16.0.0/12"),
    ipaddress.ip_network("192.168.0.0/16"),
    ipaddress.ip_network("127.0.0.0/8"),
    ipaddress.ip_network("::1/128"),
    ipaddress.ip_network("fc00::/7"),          # ULA
]

_LOOPBACK_HOSTS = {"localhost", "127.0.0.1", "::1", "0.0.0.0"}


def _configured_ai_hosts(user_id: str) -> set:
    """Hosts the user has configured as AI providers (from saved config + defaults)."""
    hosts = set()
    configs = _load_ai_configs()
    providers = configs.get(user_id) or _get_default_providers()
    for p in providers:
        base = (p.get("api_base_url") or "").strip()
        if not base:
            continue
        try:
            host = urlparse(base).hostname
            if host:
                hosts.add(host.lower())
        except Exception:
            continue
    return hosts


async def _resolve_host_ips(hostname: str):
    """Resolve hostname to IPs (non-blocking). Returns list of ipaddress objects."""
    loop = asyncio.get_event_loop()
    try:
        infos = await loop.run_in_executor(None, socket.getaddrinfo, hostname, None)
        ips = []
        for info in infos:
            try:
                ips.append(ipaddress.ip_address(info[4][0]))
            except ValueError:
                continue
        return ips
    except Exception:
        return []


async def _validate_target_url(url: str, user_id: str):
    """Validate target. Trả IP đã validate (str) để proxy connect thẳng vào đó
    chống DNS rebinding TOCTOU; None khi giữ hostname (loopback / https)."""
    parsed = urlparse(url)
    if parsed.scheme not in ("http", "https"):
        raise HTTPException(status_code=400, detail="URL chỉ hỗ trợ giao thức http/https")
    if parsed.username or parsed.password:
        raise HTTPException(status_code=400, detail="URL không được chứa thông tin đăng nhập")
    hostname = (parsed.hostname or "").lower()
    if not hostname:
        raise HTTPException(status_code=400, detail="URL không hợp lệ")

    allowed_hosts = _configured_ai_hosts(user_id)

    # Hostname-level fast path for loopback hosts
    if hostname in _LOOPBACK_HOSTS:
        if hostname in allowed_hosts:
            return None
        raise HTTPException(status_code=403, detail="Target nội bộ không nằm trong danh sách AI provider đã cấu hình")

    # Try direct IP parse (hostname may itself be an IP)
    try:
        ip = ipaddress.ip_address(hostname)
        ips = [ip]
    except ValueError:
        ips = await _resolve_host_ips(hostname)

    if not ips:
        raise HTTPException(status_code=502, detail="Không phân giải được hostname")

    for ip in ips:
        if any(ip in net for net in _BLOCKED_NETWORKS):
            raise HTTPException(status_code=403, detail="Target bị chặn (metadata/link-local không được phép)")
        if any(ip in net for net in _PRIVATE_NETWORKS):
            if hostname in allowed_hosts:
                continue
            raise HTTPException(status_code=403, detail="Target IP nội bộ không nằm trong danh sách AI provider đã cấu hình")

    # Trả IP đầu tiên đã validate — http sẽ connect thẳng vào IP này (bind),
    # không cho httpx re-resolve hostname (fix DNS rebinding TOCTOU).
    return str(ips[0])


@router.post("/proxy")
async def proxy_llm(req: ProxyRequest, current_user: dict = Depends(get_current_user)):
    target_ip = await _validate_target_url(req.url, current_user["user_id"])
    # Never forward the app's own auth token upstream.
    headers = {
        k: v for k, v in req.headers.items()
        if k.lower() not in ("host", "origin", "referer", "x-auth-token")
    }
    url = req.url
    # Bug #6: http + hostname (không phải IP literal) → connect thẳng IP đã
    # validate, giữ Host gốc. https giữ hostname (SNI + cert validation chống
    # rebinding sẵn). IPv6 skip (netloc bracket phức tạp, hiếm gặp).
    if target_ip and ":" not in target_ip:
        parsed = urlparse(req.url)
        if parsed.scheme == "http":
            host_header = parsed.netloc
            new_netloc = parsed.netloc.replace(parsed.hostname, target_ip)
            from urllib.parse import urlunsplit
            url = urlunsplit((parsed.scheme, new_netloc, parsed.path, parsed.query, parsed.fragment))
            headers["Host"] = host_header
    try:
        async with httpx.AsyncClient(timeout=180.0, follow_redirects=False) as client:
            resp = await client.post(url, headers=headers, json=req.body)
            raw = resp.text
            try:
                return resp.json()
            except json.JSONDecodeError:
                try:
                    return json.loads(raw[:raw.find('\n')])
                except (json.JSONDecodeError, ValueError):
                    return {"content": raw}
    except httpx.TimeoutException:
        raise HTTPException(status_code=504, detail="AI provider timeout")
    except httpx.ConnectError as e:
        msg = f"Cannot connect to AI provider: {e}"
        if 'localhost' in req.url or '127.0.0.1' in req.url:
            msg += "\nNếu app chạy trong Docker, localhost trỏ vào container, không ra host.\nHãy thay localhost bằng host.docker.internal hoặc IP bridge Docker (172.17.0.1)."
        raise HTTPException(status_code=502, detail=msg)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
