from typing import Annotated
from contextlib import asynccontextmanager
from importlib.resources import files
import httpx
import json
import datetime
import re

import uuid
from fastapi import FastAPI, Query, HTTPException, Request, Response, Cookie
from fastapi.responses import HTMLResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from sqlmodel import select, desc, asc
from pydantic import BaseModel

from agentbeats.model import (
    Assessment,
    create_db_and_tables,
    SessionDep,
    Agent,
    AgentDeployType,
    AgentHostedStatus,
    User,
    LogEntry,
)
from agentbeats.settings import BackendSettings
from agentbeats.myutils import check_agent


@asynccontextmanager
async def lifespan(app: FastAPI):
    create_db_and_tables()
    yield


app = FastAPI(lifespan=lifespan)
static_path = files("agentbeats.frontend").joinpath("static")
app.mount("/static", StaticFiles(directory=str(static_path)), name="static")


@app.post("/agents/")
def create_agent(
    agent: Agent, session: SessionDep, ab_api_key: str | None = Cookie(None)
) -> Agent:
    if not re.match(r"^[a-zA-Z0-9_-]+$", agent.name):
        raise HTTPException(
            status_code=400,
            detail="Agent name can only contain letters, numbers, underscores (_), and hyphens (-)",
        )

    if agent.created_at is not None:
        raise HTTPException(status_code=400, detail="created_at should not be provided")
    if agent.hosted_status != AgentHostedStatus.PENDING:
        raise HTTPException(
            status_code=400, detail="hosted_status should not be provided"
        )
    user = session.exec(select(User).where(User.ab_api_key == ab_api_key)).first()
    if agent.deploy_type != AgentDeployType.REMOTE and user.github_id not in (
        "23360163",
    ):
        raise HTTPException(
            status_code=403,
            detail="You are not authorized to create this type of agent",
        )
    if agent.deploy_type == AgentDeployType.HOSTED:
        git_fields_provided = (agent.git_url is not None) and (
            agent.git_branch is not None
        )
        docker_field_provided = agent.docker_image_url is not None
        if not (git_fields_provided or docker_field_provided):
            raise HTTPException(
                status_code=400,
                detail="For hosted agents, either git_url and git_branch or docker_image_url must be provided",
            )
    else:
        pass
    session.add(agent)
    session.commit()
    session.refresh(agent)
    return agent


@app.get("/agents/")
def read_agents(
    session: SessionDep,
    offset: int = 0,
    limit: Annotated[int, Query(ge=1, le=100)] = 100,
    role: str | None = None,  # "green" or "white" or None for all
    owner_user_id: str | None = None,
) -> list[Agent]:
    query = select(Agent)
    if role == "green":
        query = query.where(Agent.is_green)
    elif role == "white":
        query = query.where(Agent.is_green == False)
    if owner_user_id is not None:
        query = query.where(Agent.user_id == owner_user_id)
    agents = session.exec(query.offset(offset).limit(limit)).all()
    return list(agents)


@app.get("/agents/{agent_id}")
def read_agent(agent_id: str, session: SessionDep) -> Agent:
    agent = session.get(Agent, uuid.UUID(agent_id))
    if not agent:
        raise HTTPException(status_code=404, detail="Agent not found")
    return agent


@app.delete("/agents/{agent_id}")
def delete_agent(
    agent_id: str, session: SessionDep, ab_api_key: str | None = Cookie(None)
) -> dict:
    agent = session.get(Agent, uuid.UUID(agent_id))
    if not agent:
        raise HTTPException(status_code=404, detail="Agent not found")
    # check if deleted by owner
    user = session.exec(select(User).where(User.ab_api_key == ab_api_key)).first()
    assert user is not None, "User not found"
    if agent.user_id != user.github_id:
        raise HTTPException(
            status_code=403, detail="You are not authorized to delete this agent"
        )
    # first delete all linked agent checks
    for check in agent.agent_checks:
        session.delete(check)
    session.delete(agent)
    session.commit()
    return {"message": "Agent deleted"}


@app.get("/agents/{agent_id}/most_recent_check")
def get_most_recent_agent_check(agent_id: str, session: SessionDep):
    agent = session.get(Agent, uuid.UUID(agent_id))
    if not agent:
        raise HTTPException(status_code=404, detail="Agent not found")
    if not agent.agent_checks or len(agent.agent_checks) == 0:
        raise HTTPException(
            status_code=404, detail="No agent checks found for this agent"
        )
    most_recent_check = max(agent.agent_checks, key=lambda check: check.created_at)
    return most_recent_check


@app.get("/agents/{agent_id}/check_again")
def check_agent_again(agent_id: str, session: SessionDep):
    agent = session.get(Agent, uuid.UUID(agent_id))
    if not agent:
        raise HTTPException(status_code=404, detail="Agent not found")
    check_agent(session, agent)
    return {"message": "Agent check initiated"}


class AssessmentCreate(BaseModel):
    agents: list[str]
    config: str
    repeat_n: int


@app.post("/assessments/")
async def create_assessment(
    assessment_data: AssessmentCreate,
    session: SessionDep,
    ab_api_key: str | None = Cookie(None),
):
    agents = assessment_data.agents
    config = assessment_data.config
    repeat_n = assessment_data.repeat_n
    # check if all agents exist and the first is green
    if len(agents) < 2:
        raise HTTPException(
            status_code=400, detail="At least two agents (green and white) are required"
        )
    green_agent = session.get(Agent, uuid.UUID(agents[0]))
    if not green_agent or not green_agent.is_green:
        raise HTTPException(status_code=400, detail="The first agent must be green")
    for agent_id in agents[1:]:
        agent = session.get(Agent, uuid.UUID(agent_id))
        if not agent or agent.is_green:
            raise HTTPException(
                status_code=400, detail="Subsequent agents must be white agents"
            )
    # create assessment record
    assessment_ids = []
    user = session.exec(select(User).where(User.ab_api_key == ab_api_key)).first()
    assert user is not None, "User not found"
    first_assessment_id = None
    for repeat_idx in range(repeat_n):
        assessment = Assessment(
            batch_first_id=first_assessment_id,
            submitted_by=user.github_id,
            green_agent_id=agents[0],
            agent_id_list_json=json.dumps(agents),
            created_at=datetime.datetime.utcnow(),
            config=config,
            repeat_n=repeat_n,
            repeat_idx=repeat_idx,
        )
        session.add(assessment)
        session.commit()
        session.refresh(assessment)
        if first_assessment_id is None:
            first_assessment_id = assessment.id
            assessment.batch_first_id = first_assessment_id
            session.add(assessment)
            session.commit()
            session.refresh(assessment)
        assessment_ids.append(assessment.id)
        # links
        for agent_id in agents:
            agent = session.get(Agent, uuid.UUID(agent_id))
            assert agent is not None
            assessment.agents.append(agent)
        session.add(assessment)
        session.commit()
    return assessment_ids


@app.get("/assessments/{assessment_id}")
def read_assessment(assessment_id: str, session: SessionDep) -> Assessment:
    assessment = session.get(Assessment, uuid.UUID(assessment_id))
    if not assessment:
        raise HTTPException(status_code=404, detail="Assessment not found")
    return assessment


@app.delete("/assessments/{assessment_id}")
def delete_assessment(
    assessment_id: str, session: SessionDep, ab_api_key: str | None = Cookie(None)
) -> dict:
    assessment = session.get(Assessment, uuid.UUID(assessment_id))
    if not assessment:
        raise HTTPException(status_code=404, detail="Assessment not found")
    user = session.exec(select(User).where(User.ab_api_key == ab_api_key)).first()
    assert user is not None, "User not found"
    if assessment.submitted_by != user.github_id:
        raise HTTPException(
            status_code=403, detail="You are not authorized to delete this assessment"
        )
    for log in assessment.log_entries:
        session.delete(log)
    session.delete(assessment)
    session.commit()
    return {"message": "Assessment deleted"}


@app.get("/assessments/")
def read_assessments(
    session: SessionDep,
    user_id: str | None = None,
    green_agent_id: str | None = None,
    status: str | None = None,
    offset: int = 0,
    limit: Annotated[int, Query(ge=1, le=100)] = 100,
) -> list[Assessment]:
    query = select(Assessment)
    if user_id is not None:
        query = query.where(Assessment.submitted_by == user_id)
    if green_agent_id is not None:
        query = query.where(Assessment.green_agent_id == green_agent_id)
    if status is not None:
        query = query.where(Assessment.status == status)
    # sorted by created_at desc
    query = query.order_by(desc(Assessment.created_at))
    assessments = session.exec(query.offset(offset).limit(limit)).all()
    return list(assessments)


@app.get("/assessments/{assessment_id}/logs")
def read_assessment_logs(
    assessment_id: str,
    session: SessionDep,
    source: str | None = None,
    channel: str | None = None,
) -> list[LogEntry]:
    assessment = session.get(Assessment, uuid.UUID(assessment_id))
    if not assessment:
        raise HTTPException(status_code=404, detail="Assessment not found")
    query = select(LogEntry).where(LogEntry.assessment_id == uuid.UUID(assessment_id))
    if source is not None:
        query = query.where(LogEntry.source == source)
    if channel is not None:
        query = query.where(LogEntry.channel == channel)
    query = query.order_by(asc(LogEntry.created_at))
    return list(session.exec(query).all())


@app.middleware("http")
async def auth_middleware(request: Request, call_next):
    path = request.url.path
    if path != "/main":
        return await call_next(request)
    ab_api_key = request.cookies.get("ab_api_key")
    if not ab_api_key:
        return RedirectResponse(url="/login/view", status_code=303)
    return await call_next(request)


@app.get("/login/view", response_class=HTMLResponse)
async def login_view():
    html_content = files("agentbeats.frontend").joinpath("login.html").read_text()
    return html_content


@app.get("/login/callback")
async def login_callback(code: str, response: Response, session: SessionDep):
    # Exchange this code for an access token
    backend_settings = BackendSettings()
    assert backend_settings.client_secret is not None, "client_secret is not set"
    token_url = "https://github.com/login/oauth/access_token"
    data = {
        "client_id": "Ov23liEDU5GtCaixVPdg",
        "client_secret": backend_settings.client_secret,
        "code": code,
    }
    headers = {"Accept": "application/json"}
    async with httpx.AsyncClient() as client:
        token_response = await client.post(token_url, data=data, headers=headers)
        token_data = token_response.json()
        access_token = token_data.get("access_token")
        if not access_token:
            raise HTTPException(
                status_code=400, detail="Failed to retrieve access token"
            )
        # use the access token to get user info
        user_url = "https://api.github.com/user"
        auth_headers = {"Authorization": f"Bearer {access_token}"}
        user_response = await client.get(user_url, headers=auth_headers)
        user_data = user_response.json()
        user_email_url = "https://api.github.com/user/emails"
        email_response = await client.get(user_email_url, headers=auth_headers)
        email_data = email_response.json()
    # reg user or update db, and also set a cookie
    github_id = user_data["id"]
    display_name = user_data["name"] if user_data["name"] else user_data["login"]
    github_user_json = json.dumps(user_data)
    github_emails_json = json.dumps(email_data)
    user = session.get(User, github_id)
    if not user:
        # create new user
        ab_api_key = uuid.uuid4().hex
        user = User(
            github_id=github_id,
            display_name=display_name,
            github_user_json=github_user_json,
            github_emails_json=github_emails_json,
            ab_api_key=ab_api_key,
            created_at=datetime.datetime.now(),
        )
        session.add(user)
        session.commit()
        session.refresh(user)
    else:
        # update existing user
        user.display_name = display_name
        user.github_user_json = github_user_json
        user.github_emails_json = github_emails_json
        ab_api_key = user.ab_api_key
        session.add(user)
        session.commit()
        session.refresh(user)
    # set cookie
    r = RedirectResponse(url="/main")
    r.set_cookie("ab_api_key", ab_api_key, httponly=True, secure=True)
    return r


@app.get("/logout")
async def logout(response: Response):
    r = RedirectResponse(url="/login/view")
    r.delete_cookie("ab_api_key")
    return r


@app.get("/user")
async def get_user(
    session: SessionDep,
    ab_api_key: str | None = Cookie(None),
):
    if not ab_api_key:
        raise HTTPException(status_code=401, detail="Not authenticated")
    user = session.exec(select(User).where(User.ab_api_key == ab_api_key)).first()
    if not user:
        raise HTTPException(status_code=404, detail="User not found")
    return user


@app.get("/check_if_user_has_berkeley_email/{user_id}")
def check_if_user_has_berkeley_email(user_id: str, session: SessionDep) -> dict:
    user = session.get(User, user_id)
    if not user:
        raise HTTPException(status_code=404, detail="User not found")
    emails = [e["email"] for e in json.loads(user.github_emails_json) if e["verified"]]
    has_berkeley_email = any(
        email.endswith("@berkeley.edu") or email.endswith("@cs.berkeley.edu")
        for email in emails
    )
    return {"has_berkeley_email": has_berkeley_email}


@app.get("/user/renew_api_key")
async def renew_api_key(
    session: SessionDep,
    ab_api_key: str | None = Cookie(None),
):
    if not ab_api_key:
        raise HTTPException(status_code=401, detail="Not authenticated")
    user = session.exec(select(User).where(User.ab_api_key == ab_api_key)).first()
    if not user:
        raise HTTPException(status_code=404, detail="User not found")
    new_api_key = uuid.uuid4().hex
    user.ab_api_key = new_api_key
    session.add(user)
    session.commit()
    session.refresh(user)
    return {"ab_api_key": new_api_key}


@app.get("/main", response_class=HTMLResponse)
async def main_view():
    html_content = files("agentbeats.frontend").joinpath("main.html").read_text()
    return html_content


@app.get("/view/agent/{agent_id}", response_class=HTMLResponse)
async def view_agent(agent_id: str):
    html_content = files("agentbeats.frontend").joinpath("view_agent.html").read_text()
    return html_content


@app.get("/view/assessment/{assessment_id}", response_class=HTMLResponse)
async def view_assessment(assessment_id: str):
    html_content = (
        files("agentbeats.frontend").joinpath("view_assessment.html").read_text()
    )
    return html_content


@app.get("/")
async def root():
    # redirect to test_agent_crud
    return RedirectResponse(url="/main")


@app.get("/favicon.ico")
def favicon():
    return RedirectResponse(url="/static/favicon.ico")


# TEST ONLY


@app.get("/test/add_test_tau_agents")
def add_test_tau_agents(session: SessionDep) -> dict:
    """Add test Tau agents to the database for testing purposes"""

    test_user = User(
        github_id="test_user",
        display_name="Test User",
        github_user_json="{}",
        github_emails_json="[]",
        ab_api_key=uuid.uuid4().hex,
        created_at=datetime.datetime.utcnow(),
    )
    session.add(test_user)
    session.commit()
    session.refresh(test_user)
    tau_green = Agent(
        name="Tau Green",
        created_at=datetime.datetime.utcnow(),
        is_green=True,
        deploy_type=AgentDeployType.HOSTED,
        ctrl_url="",
        agent_url="",
        secret='{"ROLE": "green"}',
        inject_litellm_proxy_api=True,
        git_url="",
        git_branch="",
        docker_image_url="us-west1-docker.pkg.dev/agentbeats/gcr-test/taubench-example:latest",
        description_prompt="",
        hosted_status=AgentHostedStatus.PENDING,
        hosted_error_msg=None,
        gcp_docker_image_url=None,
        gcp_cloudrun_service_name=None,
        user_id=test_user.github_id,
    )
    tau_white = Agent(
        name="Tau White",
        created_at=datetime.datetime.utcnow(),
        is_green=False,
        deploy_type=AgentDeployType.HOSTED,
        ctrl_url="",
        agent_url="",
        secret='{"ROLE": "white"}',
        inject_litellm_proxy_api=True,
        git_url="",
        git_branch="",
        docker_image_url="us-west1-docker.pkg.dev/agentbeats/gcr-test/taubench-example:latest",
        description_prompt="",
        hosted_status=AgentHostedStatus.PENDING,
        hosted_error_msg=None,
        gcp_docker_image_url=None,
        gcp_cloudrun_service_name=None,
        user_id=test_user.github_id,
    )
    session.add(tau_green)
    session.add(tau_white)
    session.commit()
    session.refresh(tau_green)
    session.refresh(tau_white)
    return {"message": "Test Tau agents added successfully"}


@app.get("/test/view_agent_crud", response_class=HTMLResponse)
async def test_agent_crud():
    """Serve the Agent CRUD frontend interface"""
    html_content = files("agentbeats.frontend").joinpath("agent_crud.html").read_text()
    return html_content
