from fastapi import FastAPI, Body, HTTPException, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from typing import List, Dict, Any
from pathlib import Path
from datetime import datetime
import subprocess
import uuid
import pandas as pd

app = FastAPI(title="Agent Sandbox API")

class IncommingMail(BaseModel):
    data: List[Dict[str, Any]]   # requires a "data" field in JSON
    timestamp: str = datetime.utcnow().strftime("%Y%m%dT%H%M%SZ")
    id: str = uuid.uuid4().hex[:5]
    
    @property 
    def file_path(self):
        return f"{self.timestamp}_{self.id}_sample.parquet"
        
    def store_mail(self):
        try:
            import pandas as pd 
            data = pd.DataFrame(self.data)
            data.to_parquet(self.file_path)
        except Exception as e:
            raise e

class MailResponse(BaseModel):
    status: str
    timestamp: str
    pigeon_hole: str
    pigeon_location: str

class CodeScript(BaseModel):
    code: str
    
@app.post("/run")
def run_code(payload: CodeScript):
    """
    Execute Python code in a sandboxed subprocess.
    Input JSON: {"code": "print(2+2)"} or {"code": "2+2"}
    Output JSON: {"output": "4"} or {"error": "..."}
    """
    code = payload.code
    
    if not code:
        raise HTTPException(status_code=400, detail="Received Payload in incorrect format. Must pass key: 'code' with value of the script to run python code.")

    try:
        # Run the code in a subprocess with a 5s timeout
        proc = subprocess.run(
            ["python3", "-c", code],
            capture_output=True,
            text=True,
            timeout=5
        )

        if proc.returncode != 0:
            return {"error": proc.stderr.strip() or "Unknown error"}

        output = proc.stdout.strip()

        # If nothing printed, try evaluating expression directly
        if not output:
            try:
                result = eval(code, {})
                output = str(result)
            except Exception:
                output = "<no output>"

        return {"output": output}

    except subprocess.TimeoutExpired:
        raise HTTPException(status_code=408, detail="Execution timed out")

@app.post("/inbox", response_model=MailResponse)
def agent_mailbox(mail: IncommingMail):
    """
    Save incoming data with a timestamp + unique ID.
    """
    # Prepare / fetch directory
    inbox_dir = Path("inbox")
    inbox_dir.mkdir(exist_ok=True)
    mail.store_mail()
    
    return MailResponse(
        status="Data Saved!",
        timestamp=mail.timestamp,
        pigeon_hole=mail.id,
        pigeon_location=mail.file_path
    )