api_client.py

import secrets
import logging
import httpx
from datetime import date
from sqlalchemy import func

from database import SessionLocal, Provider, VirtualKey, UsageLog, Setting
from encryption import encrypt_key, decrypt_key
from export_import import export_data, import_data

logger = logging.getLogger("vapi-client")

class VApiClient:

    def get_providers(self) -> list[dict]:
        db = SessionLocal()
        try:
            rows = db.query(Provider).order_by(Provider.id).all()
            return [
                {
                    "id": p.id,
                    "name": p.name,
                    "is_active": p.is_active,
                    "created_at": p.created_at.strftime("%Y-%m-%d %H:%M"),
                }
                for p in rows
            ]
        except Exception as e:
            logger.error(f"get_providers: {e}")
            return []
        finally:
            db.close()

    def save_provider(self, name: str, api_key: str) -> bool:
        db = SessionLocal()
        try:
            existing = (
                db.query(Provider)
                .filter(func.lower(Provider.name) == func.lower(name))
                .first()
            )
            encrypted = encrypt_key(api_key)
            if existing:
                existing.api_key_encrypted = encrypted
                existing.is_active = True
            else:
                db.add(Provider(name=name, api_key_encrypted=encrypted, is_active=True))
            db.commit()
            return True
        except Exception as e:
            logger.error(f"save_provider: {e}")
            db.rollback()
            return False
        finally:
            db.close()

    def delete_provider(self, provider_id: int) -> bool:
        db = SessionLocal()
        try:
            p = db.query(Provider).filter(Provider.id == provider_id).first()
            if p:
                db.delete(p)
                db.commit()
                return True
            return False
        except Exception as e:
            logger.error(f"delete_provider: {e}")
            db.rollback()
            return False
        finally:
            db.close()

    def test_provider_connection(self, provider_id: int) -> tuple[bool, str]:
        db = SessionLocal()
        try:
            p = db.query(Provider).filter(Provider.id == provider_id).first()
            if not p:
                return False, "Provider not found"
            real_key = decrypt_key(p.api_key_encrypted)
            name = p.name.lower()
            if name == "gemini":
                url = f"https://generativelanguage.googleapis.com/v1beta/models?key={real_key}"
            elif name == "openai":
                url = "https://api.openai.com/v1/models"
            else:
                return False, f"Unknown provider type: {p.name}"
        finally:
            db.close()

        import time
        try:
            headers = {}
            if name == "openai":
                headers["Authorization"] = f"Bearer {real_key}"
            t0 = time.perf_counter()
            resp = httpx.get(url, headers=headers, timeout=10.0)
            ms = int((time.perf_counter() - t0) * 1000)
            if 200 <= resp.status_code < 300:
                return True, f"Latency: {ms}ms\nResponse: OK ({resp.status_code})"
            else:
                return False, f"Latency: {ms}ms\nResponse: HTTP {resp.status_code}"
        except Exception as e:
            return False, str(e)

    def get_virtual_keys(self) -> list[dict]:
        db = SessionLocal()
        try:
            rows = db.query(VirtualKey).order_by(VirtualKey.id).all()
            return [
                {
                    "id": vk.id,
                    "key": vk.key,
                    "name": vk.name,
                    "provider_id": vk.provider_id,
                    "provider_name": vk.provider.name if vk.provider else "Unknown",
                    "is_active": vk.is_active,
                    "daily_budget": vk.daily_budget,
                    "monthly_budget": vk.monthly_budget,
                    "allowed_models": vk.allowed_models or "",
                    "created_at": vk.created_at.strftime("%Y-%m-%d %H:%M"),
                }
                for vk in rows
            ]
        except Exception as e:
            logger.error(f"get_virtual_keys: {e}")
            return []
        finally:
            db.close()

    def create_virtual_key(self, name: str, provider_id: int, allowed_models: str = "") -> str:
        db = SessionLocal()
        try:
            token = f"vapi-sk-{secrets.token_hex(16)}"
            db.add(VirtualKey(key=token, name=name, provider_id=provider_id, is_active=True, allowed_models=allowed_models))
            db.commit()
            return token
        except Exception as e:
            logger.error(f"create_virtual_key: {e}")
            db.rollback()
            return ""
        finally:
            db.close()

    def toggle_virtual_key(self, vk_id: int, is_active: bool) -> bool:
        db = SessionLocal()
        try:
            vk = db.query(VirtualKey).filter(VirtualKey.id == vk_id).first()
            if vk:
                vk.is_active = is_active
                db.commit()
                return True
            return False
        except Exception as e:
            logger.error(f"toggle_virtual_key: {e}")
            db.rollback()
            return False
        finally:
            db.close()

    def rename_virtual_key(self, vk_id: int, new_name: str) -> bool:
        db = SessionLocal()
        try:
            vk = db.query(VirtualKey).filter(VirtualKey.id == vk_id).first()
            if vk:
                vk.name = new_name
                db.commit()
                return True
            return False
        except Exception as e:
            logger.error(f"rename_virtual_key: {e}")
            db.rollback()
            return False
        finally:
            db.close()

    def regenerate_virtual_key(self, vk_id: int) -> str:
        db = SessionLocal()
        try:
            vk = db.query(VirtualKey).filter(VirtualKey.id == vk_id).first()
            if vk:
                new_token = f"vapi-sk-{secrets.token_hex(16)}"
                vk.key = new_token
                db.commit()
                return new_token
            return ""
        except Exception as e:
            logger.error(f"regenerate_virtual_key: {e}")
            db.rollback()
            return ""
        finally:
            db.close()

    def revoke_virtual_key(self, vk_id: int) -> bool:
        db = SessionLocal()
        try:
            vk = db.query(VirtualKey).filter(VirtualKey.id == vk_id).first()
            if vk:
                db.delete(vk)
                db.commit()
                return True
            return False
        except Exception as e:
            logger.error(f"revoke_virtual_key: {e}")
            db.rollback()
            return False
        finally:
            db.close()

    def get_logs(self, limit: int = 200, **filters) -> list[dict]:
        db = SessionLocal()
        try:
            query = db.query(UsageLog)

            if "start_date" in filters and filters["start_date"]:
                query = query.filter(UsageLog.timestamp >= filters["start_date"])
            if "end_date" in filters and filters["end_date"]:
                query = query.filter(UsageLog.timestamp <= filters["end_date"])
            if "status_code" in filters and filters["status_code"]:
                sc = filters["status_code"]
                if sc == "Success (2xx)":
                    query = query.filter(UsageLog.status_code >= 200, UsageLog.status_code < 300)
                elif sc == "Error (4xx/5xx)":
                    query = query.filter(UsageLog.status_code >= 400)
                else:
                    try:
                        query = query.filter(UsageLog.status_code == int(sc))
                    except ValueError:
                        pass
            if "virtual_key_id" in filters and filters["virtual_key_id"]:
                query = query.filter(UsageLog.virtual_key_id == filters["virtual_key_id"])
            if "provider_id" in filters and filters["provider_id"]:
                query = query.join(VirtualKey).filter(VirtualKey.provider_id == filters["provider_id"])

            rows = query.order_by(UsageLog.timestamp.desc()).limit(limit).all()
            return [
                {
                    "id": log.id,
                    "virtual_key_name": log.virtual_key.name if log.virtual_key else "Revoked",
                    "virtual_key": log.virtual_key.key if log.virtual_key else "N/A",
                    "endpoint": log.endpoint,
                    "method": log.method,
                    "status_code": log.status_code,
                    "cost": log.cost,
                    "timestamp": log.timestamp.strftime("%Y-%m-%d %H:%M:%S"),
                }
                for log in rows
            ]
        except Exception as e:
            logger.error(f"get_logs: {e}")
            return []
        finally:
            db.close()

    def get_usage_stats(self) -> dict:
        db = SessionLocal()
        try:
            total = db.query(UsageLog).count()
            success = db.query(UsageLog).filter(
                UsageLog.status_code >= 200, UsageLog.status_code < 300
            ).count()
            failed = db.query(UsageLog).filter(UsageLog.status_code >= 400).count()

            today = date.today().strftime("%Y-%m-%d")
            today_count = db.query(UsageLog).filter(
                func.date(UsageLog.timestamp) == today
            ).count()

            providers_count = db.query(Provider).count()
            keys_count = db.query(VirtualKey).count()

            return {
                "total_requests": total,
                "success_requests": success,
                "failed_requests": failed,
                "requests_today": today_count,
                "providers_count": providers_count,
                "keys_count": keys_count,
            }
        except Exception as e:
            logger.error(f"get_usage_stats: {e}")
            return {
                "total_requests": 0, "success_requests": 0,
                "failed_requests": 0, "requests_today": 0,
                "providers_count": 0, "keys_count": 0,
            }
        finally:
            db.close()

    def get_setting(self, key: str, default: str = "") -> str:
        db = SessionLocal()
        try:
            s = db.query(Setting).filter(Setting.key == key).first()
            return s.value if s else default
        except Exception:
            return default
        finally:
            db.close()

    def save_setting(self, key: str, value: str) -> bool:
        db = SessionLocal()
        try:
            s = db.query(Setting).filter(Setting.key == key).first()
            if s:
                s.value = value
            else:
                db.add(Setting(key=key, value=value))
            db.commit()
            return True
        except Exception as e:
            logger.error(f"save_setting: {e}")
            db.rollback()
            return False
        finally:
            db.close()

    def export_config(self, path: str) -> bool:
        return export_data(path)

    def import_config(self, path: str) -> bool:
        return import_data(path)

    def run_diagnostics(self) -> list[dict]:
        results = []
        
        db_ok = False
        try:
            db = SessionLocal()
            db.query(Provider).first()
            db_ok = True
            db.close()
        except Exception:
            pass
        results.append({"name": "Database", "status": "🟢" if db_ok else "🔴", "latency": "N/A"})
        
        enc_ok = False
        try:
            test_str = "vapi_test"
            enc = encrypt_key(test_str)
            dec = decrypt_key(enc)
            enc_ok = (dec == test_str)
        except Exception:
            pass
        results.append({"name": "Encryption", "status": "🟢" if enc_ok else "🔴", "latency": "N/A"})
        
        int_ok = False
        import time
        t0 = time.perf_counter()
        try:
            resp = httpx.get("https://1.1.1.1", timeout=5.0)
            int_ok = True
            int_ms = int((time.perf_counter() - t0) * 1000)
        except Exception:
            int_ms = 0
        results.append({"name": "Internet Access", "status": "🟢" if int_ok else "🔴", "latency": f"{int_ms}ms" if int_ms else "Error"})
        
        gem_ok = False
        t0 = time.perf_counter()
        try:
            resp = httpx.get("https://generativelanguage.googleapis.com", timeout=5.0)
            gem_ok = True
            gem_ms = int((time.perf_counter() - t0) * 1000)
        except Exception:
            gem_ms = 0
        results.append({"name": "Gemini Reachability", "status": "🟢" if gem_ok else "🔴", "latency": f"{gem_ms}ms" if gem_ms else "Error"})
        
        oai_ok = False
        t0 = time.perf_counter()
        try:
            resp = httpx.get("https://api.openai.com/v1/", timeout=5.0)
            oai_ok = True
            oai_ms = int((time.perf_counter() - t0) * 1000)
        except Exception:
            oai_ms = 0
        results.append({"name": "OpenAI Reachability", "status": "🟢" if oai_ok else "🔴", "latency": f"{oai_ms}ms" if oai_ms else "Error"})
        
        gw_ok = self.check_gateway_health()
        results.append({"name": "Gateway Thread", "status": "🟢" if gw_ok else "🔴", "latency": "Local"})
        
        return results

    def check_gateway_health(self, port: int = 8000) -> bool:
        try:
            resp = httpx.get(f"http://127.0.0.1:{port}/health", timeout=2.0)
            return resp.status_code == 200
        except Exception:
            return False