import os
import psutil
from functools import lru_cache
#
import cpuinfo
from fastapi import FastAPI, Request, Body
from coreference_resolution import CoreferenceResolver

app = FastAPI()

@lru_cache(maxsize=1)
def get_resolver():
    return CoreferenceResolver()

resolver = get_resolver()

@app.post(path="/api/v1/fastcoref")
def get_coref_resolution(text:str = Body(...)) -> dict:
    result = resolver.resolve_with_confidence(text)
    return result
    
def get_container_cpu_count():
    try:
        # default CPU count from Python
        cpu_count = os.cpu_count()

        # read cgroup quota (if available)
        quota_path = "/sys/fs/cgroup/cpu/cpu.cfs_quota_us"
        period_path = "/sys/fs/cgroup/cpu/cpu.cfs_period_us"
        if os.path.exists(quota_path) and os.path.exists(period_path):
            with open(quota_path) as f:
                quota = int(f.read())
            with open(period_path) as f:
                period = int(f.read())
            if quota > 0 and period > 0:
                cpu_count = quota / period
        # else:
        #     print(f"{quota_path} or {period_path} doesn't exist")
        return int(cpu_count)
    except Exception:
        return os.cpu_count()

def get_container_memory_limit():
    try:
        mem_limit_path = "/sys/fs/cgroup/memory/memory.limit_in_bytes"
        if os.path.exists(mem_limit_path):
            with open(mem_limit_path) as f:
                mem_bytes = int(f.read())
            return mem_bytes // (1024 * 1024)  # in MB
        # else:
        #     print(f"{mem_limit_path} doesn't exist")
        # fallback to total system memory
        return int(psutil.virtual_memory().total // (1024 * 1024))
    except Exception:
        return None

# = for cgroup v2
def get_cgroupv2_cpu_limit():
    try:
        with open("/sys/fs/cgroup/cpu.max") as f:
            data = f.read().strip()
        quota, period = data.split()
        if quota == "max":
            print("no CPU limit")
            return os.cpu_count()  # no limit
        return int(int(quota) / int(period))
    except Exception as ex:
        print(repr(ex))
        return os.cpu_count()

def get_cgroupv2_memory_limit():
    try:
        with open("/sys/fs/cgroup/memory.max") as f:
            mem_bytes = f.read().strip()
        if mem_bytes == "max":
            print("no MEMORY limit")
            return int(psutil.virtual_memory().total // (1024*1024)) # no limit
        return int(mem_bytes) // (1024*1024)  # in MB
    except Exception as ex:
        print(repr(ex))
        return int(psutil.virtual_memory().total // (1024*1024))
        
@app.get("/")
def greet_json():
    return {"Hello": "World!", "cpu_count": get_cgroupv2_cpu_limit(), "total_ram_in_MB": get_cgroupv2_memory_limit(), "cpu_name":cpuinfo.get_cpu_info()['brand_raw']}

@app.get("/health")
def health():
    return "ok"