# app.py
import os
import uuid
import boto3
import supabase
import numpy as np
import cv2
import base64
import mimetypes
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from botocore.exceptions import NoCredentialsError

# ==============================================================================
# 1. CONFIGURAÇÃO E CLIENTES (CARREGADOS NA INICIALIZAÇÃO)
# ==============================================================================

# Carrega as configurações das variáveis de ambiente (Secrets do Hugging Face)
AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
AWS_S3_BUCKET_NAME = os.getenv("AWS_S3_BUCKET_NAME")
AWS_S3_REGION = os.getenv("AWS_S3_REGION")
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")

# Validação para garantir que todos os secrets foram configurados
if not all([AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_S3_BUCKET_NAME, AWS_S3_REGION, SUPABASE_URL, SUPABASE_KEY]):
    raise RuntimeError("Erro: Nem todos os secrets necessários foram configurados no Hugging Face Space.")

# Inicializa os clientes globais
try:
    s3_client = boto3.client(
        's3',
        aws_access_key_id=AWS_ACCESS_KEY_ID,
        aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
        region_name=AWS_S3_REGION
    )
    supabase_client = supabase.create_client(SUPABASE_URL, SUPABASE_KEY)
    print("Clientes S3 e Supabase inicializados com sucesso.")
except Exception as e:
    raise RuntimeError(f"Erro ao inicializar os clientes: {e}")

# Inicializa a aplicação FastAPI
app = FastAPI(title="Detector de Corrosão Branca com Salvamento S3/Supabase")

# Em produção, restrinja ao domínio do seu frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ==============================================================================
# 2. FUNÇÕES HELPER (LÓGICA REUTILIZÁVEL)
# ==============================================================================

def process_image_bytes(img_bytes: bytes):
    """Lógica de análise de imagem com OpenCV. Retorna um dict com métricas e a imagem de corrosão em numpy (RGB)."""
    nparr = np.frombuffer(img_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError("Não foi possível decodificar a imagem.")

    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    # Segmentação do objeto principal (fundo escuro)
    lower_bg = np.array([0, 0, 0], dtype=np.uint8)
    upper_bg = np.array([180, 255, 50], dtype=np.uint8)
    mask_bg = cv2.inRange(hsv, lower_bg, upper_bg)
    mask_obj = cv2.bitwise_not(mask_bg)

    kernel = np.ones((5, 5), np.uint8)
    mask_obj = cv2.morphologyEx(mask_obj, cv2.MORPH_OPEN, kernel)
    mask_obj = cv2.morphologyEx(mask_obj, cv2.MORPH_CLOSE, kernel)

    contours, _ = cv2.findContours(mask_obj, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        raise ValueError("Nenhum objeto principal detectado na imagem.")

    largest = max(contours, key=cv2.contourArea)
    mask_clean = np.zeros_like(mask_obj)
    cv2.drawContours(mask_clean, [largest], -1, 255, cv2.FILLED)

    isolated = cv2.bitwise_and(img, img, mask=mask_clean)
    hsv_iso = cv2.cvtColor(isolated, cv2.COLOR_BGR2HSV)

    # Corrosão branca (faixa de brancos)
    lower_white = np.array([0, 0, 180], dtype=np.uint8)
    upper_white = np.array([180, 60, 255], dtype=np.uint8)
    mask_white = cv2.inRange(hsv_iso, lower_white, upper_white)
    mask_white = cv2.bitwise_and(mask_white, mask_white, mask=mask_clean)

    total_pixels = int(np.count_nonzero(mask_clean))
    corrosion_pixels = int(np.count_nonzero(mask_white))
    percent = (corrosion_pixels / max(1, total_pixels)) * 100.0

    isolated_rgb = cv2.cvtColor(isolated, cv2.COLOR_BGR2RGB)
    corrosion_vis_rgb = cv2.bitwise_and(isolated_rgb, isolated_rgb, mask=mask_white)

    def to_data_uri(img_arr_rgb):
        bgr = cv2.cvtColor(img_arr_rgb, cv2.COLOR_RGB2BGR)
        ok, buf = cv2.imencode(".png", bgr)
        if not ok:
            return None
        b64 = base64.b64encode(buf.tobytes()).decode("ascii")
        return f"data:image/png;base64,{b64}"

    analysis_results = {
        "percent": round(percent, 4),
        "total_pixels": total_pixels,
        "corrosion_pixels": corrosion_pixels,
        "isolated_image": to_data_uri(isolated_rgb),
        "corrosion_image": to_data_uri(corrosion_vis_rgb),
    }

        return analysis_results, corrosion_vis_rgb

    
def analyze_and_persist(img_bytes: bytes, logical_name: str):
    """
    Processa a imagem em memória (process_image_bytes), faz upload no S3 (original & resultado)
    e insere na tabela 'amostras'. Retorna o dict da análise + chaves S3 + database_id.
    """
    # 1) Processa
    analysis_results, corrosion_image_np = process_image_bytes(img_bytes)

    # 2) Upload ORIGINAL
    content_type = mimetypes.guess_type(logical_name)[0] or "application/octet-stream"
    extensao = mimetypes.guess_extension(content_type) or ".jpg"

    s3_key_original = f"imagens_originais/{uuid.uuid4()}{extensao}"
    s3_client.put_object(
        Bucket=AWS_S3_BUCKET_NAME,
        Key=s3_key_original,
        Body=img_bytes,
        ContentType=content_type
    )

    # 3) Upload RESULTADO (PNG)
    bgr_result = cv2.cvtColor(corrosion_image_np, cv2.COLOR_RGB2BGR)
    ok, buffer = cv2.imencode('.png', bgr_result)
    if not ok:
        raise ValueError("Falha ao codificar a imagem de resultado para PNG.")
    s3_key_resultado = f"imagens_resultados/{uuid.uuid4()}.png"
    s3_client.put_object(
        Bucket=AWS_S3_BUCKET_NAME,
        Key=s3_key_resultado,
        Body=buffer.tobytes(),
        ContentType='image/png'
    )

    # 4) Insert Supabase (na mesma tabela 'amostras')
    dados_para_inserir = {
        "imagem_original": s3_key_original,
        "imagem_resultado": s3_key_resultado,
        "resultado": analysis_results["percent"]
    }
    response = supabase_client.from_("amostra").insert(dados_para_inserir).execute()
    new_record_id = response.data[0]['id']

    # Retorno enriquecido
    out = {
        **analysis_results,
        "database_id": new_record_id,
        "s3_original": s3_key_original,
        "s3_resultado": s3_key_resultado
    }
    return out

    def segmentate(file_bytes: bytes):
    """
    Roda o YOLO, salva predição e crops, e retorna:
    - save_dir (Path): diretório raiz onde o Ultralytics salvou esta execução
    - crop_paths (list[Path]): lista com todos os arquivos de imagem dentro de /crops
    """
    modelo = YOLO('back-ciser/best.pt')

    # Passa a imagem como array (sem salvar a original só para inferir)
    img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
    np_img = np.array(img)
    np_img_bgr = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)

    resultados = modelo.predict(
        source=np_img,
        imgsz=1280,
        conf=0.7,
        iou=0.3,
        save=True,
        save_crop=True,
        project=LOCAL_RESULT,
        name="predict",      # deixa previsível
        exist_ok=True,
        verbose=False
    )

    # Diretório real onde esta execução foi salva
    save_dir = Path(resultados[0].save_dir)

    crops_dir = save_dir / "crops"
    if not crops_dir.exists():
        return save_dir, []

    # Varre recursivamente (crops/<classe>/*.jpg|png|...)
    crop_paths = []
    for ext in ("*.jpg", "*.png", "*.jpeg", "*.bmp", "*.tif", "*.tiff"):
        crop_paths.extend(crops_dir.rglob(ext))

    crop_paths = sorted(crop_paths, key=lambda p: str(p).lower())
    return save_dir, crop_paths

# ==============================================================================
# 3. ENDPOINTS DA API
# ==============================================================================

@app.get("/")
def read_root():
    return {"status": "ok", "message": "API de Análise de Corrosão no ar!"}

@app.post("/analyze")
async def analyze(file: UploadFile = File(...)):
    """
    Endpoint principal: Analisa a imagem, salva original e resultado no S3,
    e persiste os dados no Supabase.
    """
    content = await file.read()
    s3_key_original = None
    s3_key_resultado = None

    try:
        # 1. Processa a imagem para obter métricas e a imagem de resultado
        analysis_results, corrosion_image_np = process_image_bytes(content)

        # 2. Faz o upload da IMAGEM ORIGINAL para o S3
        content_type = file.content_type or 'application/octet-stream'
        extensao = mimetypes.guess_extension(content_type) or '.jpg'
        s3_key_original = f"imagens_originais/{uuid.uuid4()}{extensao}"
        
        print(f"Fazendo upload da imagem original para: {s3_key_original}")
        s3_client.put_object(
            Bucket=AWS_S3_BUCKET_NAME,
            Key=s3_key_original,
            Body=content,
            ContentType=content_type
        )

        # 3. Faz o upload da IMAGEM DE RESULTADO (em memória) para o S3
        s3_key_resultado = f"imagens_resultados/{uuid.uuid4()}.png"
        
        print(f"Fazendo upload da imagem de resultado para: {s3_key_resultado}")
        bgr_result = cv2.cvtColor(corrosion_image_np, cv2.COLOR_RGB2BGR)
        ok, buffer = cv2.imencode('.png', bgr_result)
        if not ok:
            raise ValueError("Falha ao codificar a imagem de resultado para PNG.")
        
        s3_client.put_object(
            Bucket=AWS_S3_BUCKET_NAME,
            Key=s3_key_resultado,
            Body=buffer.tobytes(),
            ContentType='image/png'
        )

        # 4. Salva os metadados e as chaves S3 no Supabase
        print("Inserindo registro no Supabase...")
        dados_para_inserir = {
            #"nome_amostra": file.filename,
            #"percentual_corrosao": analysis_results["percent"],
            #"pixels_totais_obj": analysis_results["total_pixels"],
            #"pixels_corrosao": analysis_results["corrosion_pixels"],
            "imagem_original": s3_key_original,
            "imagem_resultado": s3_key_resultado,
            "resultado": analysis_results["percent"]
        }
        response = supabase_client.from_("amostra").insert(dados_para_inserir).execute()
        
        # Adiciona o ID do novo registro à resposta
        new_record_id = response.data[0]['id']
        analysis_results['database_id'] = new_record_id
        
        return JSONResponse(content=analysis_results)

    except Exception as e:
        # Lógica de Rollback: Se algo falhar, tenta limpar os arquivos do S3
        print(f"ERRO no processo de análise: {e}")
        if s3_key_original:
            print(f"Removendo objeto órfão do S3: {s3_key_original}")
            s3_client.delete_object(Bucket=AWS_S3_BUCKET_NAME, Key=s3_key_original)
        if s3_key_resultado:
            print(f"Removendo objeto órfão do S3: {s3_key_resultado}")
            s3_client.delete_object(Bucket=AWS_S3_BUCKET_NAME, Key=s3_key_resultado)
        
        raise HTTPException(status_code=500, detail=f"Ocorreu um erro interno: {e}")


@app.get("/samples/{sample_id}")
async def get_sample_images(sample_id: int):
    """
    Endpoint de recuperação: Busca uma amostra pelo ID e retorna
    links temporários (pré-assinados) para suas imagens no S3.
    """
    try:
        print(f"Buscando amostra com ID: {sample_id}...")
        response = supabase_client.from_("amostra").select("imagem_original, imagem_resultado").eq("id", sample_id).single().execute()
        
        if not response.data:
            raise HTTPException(status_code=404, detail=f"Amostra com ID {sample_id} não encontrada.")
            
        amostra = response.data
        s3_key_original = amostra.get("imagem_original")
        s3_key_resultado = amostra.get("imagem_resultado")
        
        links = {}
        
        if s3_key_original:
            links['url_original'] = s3_client.generate_presigned_url(
                'get_object',
                Params={'Bucket': AWS_S3_BUCKET_NAME, 'Key': s3_key_original},
                ExpiresIn=3600 # 1 hora
            )
        
        if s3_key_resultado:
            links['url_resultado'] = s3_client.generate_presigned_url(
                'get_object',
                Params={'Bucket': AWS_S3_BUCKET_NAME, 'Key': s3_key_resultado},
                ExpiresIn=3600 # 1 hora
            )
            
        return JSONResponse(content=links)

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Ocorreu um erro interno ao buscar os links: {e}")