import datetime
import decimal
import logging
import json
import random
import secrets
import string
import uuid
from typing import Optional, List
import time
import sys
import traceback
import dspy
import sqlalchemy
from fastapi import (
    FastAPI,
    BackgroundTasks,
    HTTPException,
    Depends,
    Header,
    status,
)
from pydantic import BaseModel
from dotenv import load_dotenv
from pydantic_core import MultiHostUrl
from sqlalchemy import create_engine, text
from sqlalchemy.exc import NoResultFound, OperationalError, SQLAlchemyError
from sqlmodel import Session
from starlette.middleware.cors import CORSMiddleware
from urllib.parse import quote_plus
from tenacity import retry, stop_after_attempt, wait_exponential

from app.core.sql_extraction import SQlExtractor
from app.db import Scoped_Session
from app.models import DBInstance, InstanceStatus
from app.models.user_session import UserSession
from app.models.validation_job import (
    ValidationJobStatus,
    QueryRun,
    Step,
    ValidationJob,
)
from app.models.regression_dataset import RegressionSample
from app.repositories.db_instance import DbInstanceRepo
from app.repositories.user_session import UserSessionRepo
from app.repositories.validation_job import ValidationJobRepo
from app.repositories.regression_dataset import RegressionDatasetRepo
from app.settings import settings
from app.core.sql_result_verifier import SQLVerifier
from app.utils.tidb_serverless import TiDBServerlessClient
from app.task_queue import task_queue
from app.simple_cache import simple_semantic_cache, extract_question

load_dotenv()

app = FastAPI()

# initial dspy programs
turbo = dspy.OpenAI(model="gpt-4o", api_key=settings.OPENAI_API_KEY, max_tokens=4096)
dspy.settings.configure(lm=turbo)

sql_result_verifier = SQLVerifier(turbo)
sql_extractor = SQlExtractor(turbo, "dspy_utils/sql_sample_extraction_program")

if settings.BACKEND_CORS_ORIGINS:
    app.add_middleware(
        CORSMiddleware,
        allow_origins=[
            str(origin).strip("/") for origin in settings.BACKEND_CORS_ORIGINS
        ],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

db_instance_repo = DbInstanceRepo(
    TiDBServerlessClient(
        public_key=settings.TiDBCLOUD_PUBLIC_KEY,
        private_key=settings.TiDBCLOUD_PRIVATE_KEY,
        project_id=settings.TiDBCLOUD_PROJECT_ID,
        cluster_id=settings.TiDBCLOUD_CLUSTER_ID,
    )
)
user_session_repo = UserSessionRepo()
validation_job_repo = ValidationJobRepo()
regression_dataset_repo = RegressionDatasetRepo()

logging.basicConfig(
    level=settings.LOG_LEVEL,
    format="[%(filename)s:%(lineno)d | %(funcName)s()] - %(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    stream=sys.stdout,  # This ensures logs go to stdout
)
logger = logging.getLogger(__name__)


def verify_token(authorization: str = Header(...)):
    if not authorization.startswith("Bearer "):
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Invalid authorization scheme",
        )

    token = authorization.split(" ")[1]
    if len(settings.STATIC_TOKEN) > 0 and token != settings.STATIC_TOKEN:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Invalid token",
        )


# start task queue workers when the app starts
@app.on_event("startup")
async def startup_event():
    task_queue.start_workers(validation_worker)


class SQLValidationRequest(BaseModel):
    external_request_id: Optional[str] = None
    qa_content: str


class CreateSQLValidationJobResponse(BaseModel):
    job_id: uuid.UUID


class QueryRequest(BaseModel):
    session_id: Optional[str] = None
    steps: list[Step]
    fast_fail: Optional[bool] = True


class QueryResult(BaseModel):
    success: bool
    session_id: uuid.UUID
    instance_id: Optional[int] = None
    message: Optional[str] = None
    runs: Optional[list[Optional[QueryRun]]] = []


class ValidationResult(BaseModel):
    status: ValidationJobStatus
    message: Optional[str]
    runs: Optional[list[Optional[QueryRun]]]


class ValidationResultReport(BaseModel):
    status: ValidationJobStatus
    message: Optional[str]
    runs_report: Optional[str]


class FailedJobsResponse(BaseModel):
    jobs: List[ValidationJob]


class ReplayJobResponse(BaseModel):
    job: ValidationJob


class AddFailedJobToDatasetResponse(BaseModel):
    success: bool
    case_id: int
    job_id: uuid.UUID


class RegressionDatasetResponse(BaseModel):
    samples: List[RegressionSample]


@app.post("/api/v1/sqls-validation")
async def create_sqls_validation_job(
    job_request: SQLValidationRequest, background_tasks: BackgroundTasks
) -> CreateSQLValidationJobResponse:
    with Scoped_Session() as manager_session:
        # try to get from semantic cache
        job_id = simple_semantic_cache.get(extract_question(job_request.qa_content))
        if job_id != None:
            return CreateSQLValidationJobResponse(job_id=job_id)

        job = None
        if job_request.external_request_id:
            job = validation_job_repo.get_validation_job_by_external_request_id(
                manager_session, job_request.external_request_id
            )
            if job and job.status != ValidationJobStatus.EXPIRED:
                return CreateSQLValidationJobResponse(job_id=job.id)

        if job is None:
            job = validation_job_repo.create_validation_job(
                manager_session, job_request.model_dump()
            )
        elif job.status == ValidationJobStatus.EXPIRED:
            # If the job is expired, we need to reset its status to CREATED and update the started_at time.
            job.status = ValidationJobStatus.CREATED
            job.started_at = datetime.datetime.now()
            manager_session.add(job)

        manager_session.commit()
        job_id = job.id

        # Schedule the validation job to run in the background.
        task_queue.add_task(job_id, job_request.model_dump())

        return CreateSQLValidationJobResponse(job_id=job_id)


@app.get("/api/v1/sqls-validation/{job_id}")
async def get_sqls_validation_job(job_id: uuid.UUID) -> ValidationResult:
    with Scoped_Session() as manager_session:
        job = validation_job_repo.get_validation_job_by_id(manager_session, job_id)
        if job is None:
            raise HTTPException(status_code=404, detail="Validation job not found.")

        return ValidationResult(
            status=job.status,
            message=job.message,
            runs=(
                [QueryRun.model_validate(run) for run in job.query_runs]
                if job.query_runs
                else []
            ),
        )


@app.post("/api/v1/jobs")
async def create_validation_job(
    job_request: SQLValidationRequest,
    background_tasks: BackgroundTasks,
    token: str = Depends(verify_token),
) -> CreateSQLValidationJobResponse:
    with Scoped_Session() as manager_session:
        # try to get from semantic cache
        job_id = simple_semantic_cache.get(extract_question(job_request.qa_content))
        if job_id != None:
            return CreateSQLValidationJobResponse(job_id=job_id)

        job = None
        if job_request.external_request_id:
            job = validation_job_repo.get_validation_job_by_external_request_id(
                manager_session, job_request.external_request_id
            )
            if job and job.status != ValidationJobStatus.EXPIRED:
                return CreateSQLValidationJobResponse(job_id=job.id)

        if job is None:
            job = validation_job_repo.create_validation_job(
                manager_session, job_request.model_dump()
            )
        elif job.status == ValidationJobStatus.EXPIRED:
            # If the job is expired, we need to reset its status to CREATED and update the started_at time.
            job.status = ValidationJobStatus.CREATED
            job.started_at = datetime.datetime.now()
            manager_session.add(job)

        manager_session.commit()
        job_id = job.id

        # Schedule the validation job to run in the background.
        task_queue.add_task(job_id, job_request.model_dump())

        return CreateSQLValidationJobResponse(job_id=job_id)


@app.get("/api/v1/jobs/{id}")
async def get_validation_job(id: str) -> ValidationResultReport:
    with Scoped_Session() as manager_session:
        job = None
        try:
            job_id = uuid.UUID(id)
            logger.info(f"Getting validation job by job id: {job_id}")
            job = validation_job_repo.get_validation_job_by_id(manager_session, job_id)
        except ValueError:
            logger.info(f"Getting validation job by external request id: {id}")
            job = validation_job_repo.get_validation_job_by_external_request_id(
                manager_session, id
            )

        if job is None:
            raise HTTPException(status_code=404, detail="Validation job not found.")

        runs = (
            [QueryRun.model_validate(run) for run in job.query_runs]
            if job.query_runs
            else []
        )
        runs_report = None
        if len(runs) > 0:
            try:
                runs_report = render_runs_to_markdown(runs)
            except Exception as e:
                logger.error(f"Failed to render runs to markdown: {e}")
                runs_report = None

        return ValidationResultReport(
            status=job.status,
            message=job.message,
            runs_report=runs_report,
        )


@app.get("/api/v1/failed-jobs")
async def get_failed_jobs(
    start_date: Optional[datetime.datetime] = None,
    end_date: Optional[datetime.datetime] = None,
) -> FailedJobsResponse:
    with Scoped_Session() as session:
        if not start_date:
            start_date = datetime.datetime.now() - datetime.timedelta(days=2)
        if not end_date:
            end_date = datetime.datetime.now()
        logger.info(f"Getting failed jobs from {start_date} to {end_date}")
        failed_jobs = ValidationJobRepo.get_failed_jobs(session, start_date, end_date)
        return FailedJobsResponse(jobs=failed_jobs)


@app.post("/api/v1/replay-job/{job_id}")
async def replay_job(job_id: uuid.UUID) -> ReplayJobResponse:
    with Scoped_Session() as session:
        job = validation_job_repo.replay_job(session, job_id)
        # Schedule the validation job to run in the background.
        task_queue.add_task(job.id, job.request)
        return ReplayJobResponse(job=job)


@app.post("/api/v1/add-job-to-dataset/{job_id}")
async def add_job_to_dataset(job_id: uuid.UUID) -> AddFailedJobToDatasetResponse:
    with Scoped_Session() as session:
        job = validation_job_repo.get_validation_job_by_id(session, job_id)
        if job is None:
            raise HTTPException(status_code=404, detail="Validation job not found.")

        try:
            sample = RegressionDatasetRepo.create_from_job(session, job)
            return AddFailedJobToDatasetResponse(
                success=True, job_id=job.id, case_id=sample.id
            )
        except Exception as e:
            logger.error(f"Failed to add job to dataset: {e}")
            return AddFailedJobToDatasetResponse(
                success=False,
                job_id=job.id,
            )


@app.get("/api/v1/regression-dataset")
async def get_regression_dataset(
    ids: Optional[str] = None, limit: int = 10, offset: int = 0
) -> RegressionDatasetResponse:
    with Scoped_Session() as session:
        if ids:
            # Convert comma-separated string to list of integers
            id_list = [int(id.strip()) for id in ids.split(",")]
            samples = RegressionDatasetRepo.get_samples(session, id_list)
        else:
            # If no specific case_ids are provided, get paginated samples
            samples = RegressionDatasetRepo.get_samples_paginated(
                session, limit, offset
            )

        if not samples:
            raise HTTPException(
                status_code=404, detail="No regression dataset samples found."
            )

        return RegressionDatasetResponse(samples=samples)


# @app.post("/api/query")
# async def query(query_request: QueryRequest):
#     return run_query(query_request)


async def validation_worker(job_id: uuid.UUID, job_request: dict):
    logger.info(f"Starting validation job: {job_id}")
    result = None
    with Scoped_Session() as manager_session:
        try:
            job = validation_job_repo.get_validation_job_by_id(manager_session, job_id)
            if job is None:
                return

            # Check if the job has expired (created more than 20 minutes ago)
            if datetime.datetime.utcnow() - job.started_at >= datetime.timedelta(
                seconds=settings.TASK_QUEUE_TIMEOUT
            ):
                job.status = ValidationJobStatus.EXPIRED
                job.message = (
                    "Job expired. It was not processed within 20 minutes of creation."
                )
                manager_session.add(job)
                manager_session.commit()
                return

            # Extract SQLs from the QA content.
            job.status = ValidationJobStatus.EXTRACTING
            job.message = "Extracting and constructing SQL examples from answer ..."
            manager_session.add(job)
            manager_session.commit()

            samples = sql_extractor.gen(job_request["qa_content"])
            if samples.no_reasonable_example or len(samples.steps) == 0:
                job.status = ValidationJobStatus.SKIPPED
                job.message = "No SQL examples found to validate."
                manager_session.add(job)
                manager_session.commit()
                return

            steps = samples.steps

            # Run the SQLs against the database.
            job.status = ValidationJobStatus.VALIDATING
            job.message = "Validating SQL statements on TiDB Serverless ..."
            job.steps = [step.model_dump() for step in steps]
            manager_session.add(job)
            manager_session.commit()

            result = run_query(QueryRequest(steps=steps, fast_fail=True))
            job = manager_session.merge(job)

            # Update the job status based on the query result.
            if result.success:
                job.status = ValidationJobStatus.SUCCESS
                job.message = "All SQL examples have been validated on TiDB Serverless."
            else:
                job.status = ValidationJobStatus.FAILED
                job.message = (
                    result.message
                    if result.message is not None
                    else "Some SQL examples did not pass validation on TiDB Serverless."
                )
            job.query_runs = [run.model_dump() for run in result.runs]
            job.query_runs = serial_helper(job.query_runs)
            job.instance_id = result.instance_id
            logger.debug(
                f"Validation job ({job_id}) completed with status: {job.status}"
            )
            manager_session.add(job)
            manager_session.commit()
        except Exception as e:
            job = manager_session.merge(job)
            logger.error("Failed to validate job (%s): %s", job.id, e)
            job.query_runs = []
            try:
                job.status = ValidationJobStatus.FAILED
                job.message = str(e)
                job.instance_id = result.instance_id if result is not None else None
                manager_session.add(job)
                manager_session.commit()
            except Exception as e:
                logger.error("Failed to update job (%s) status: %s", job.id, e)


@retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=5, min=5, max=30))
def get_db_instance_by_random(session: Session):
    try:
        instances = db_instance_repo.get_all_available_db_instance(session)
    except Exception as e:
        logger.error(f"Failed to query available instance objects from database: {e}")
        raise e

    if not instances:
        raise NotDBInstanceAvailable("No available database instances found.")

    return random.choice(instances), len(instances)


def run_query(query_request: QueryRequest) -> QueryResult:
    user_session_id = query_request.session_id or uuid.uuid4()
    instance_count = 0
    instance_id = None
    try:
        with Scoped_Session() as manager_session:
            user_session = user_session_repo.find_or_create_user_session(
                manager_session, user_session_id
            )
            if user_session.db_host is None:
                try:
                    instance, instance_count = get_db_instance_by_random(
                        manager_session
                    )
                    instance_id = instance.id
                    user_session = provision_database_for_session(
                        manager_session, instance, user_session
                    )
                    manager_session.commit()
                except NotDBInstanceAvailable as e:
                    manager_session.rollback()
                    logger.error(f"No available database instances: {e}")
                    return QueryResult(
                        success=False,
                        session_id=user_session_id,
                        message=f"No available database instances. Please try again later. {e}",
                        instance_id=instance_id,
                    )
                except OperationalError as e:
                    manager_session.rollback()
                    error_trace = traceback.format_exc()
                    logger.error(
                        f"Failed to connect to database: {e}\nStack trace:\n{error_trace}"
                    )
                    return QueryResult(
                        success=False,
                        session_id=user_session_id,
                        message=f"Failed to connect to database. Please try again later. {e}",
                        instance_id=instance_id,
                    )
                except Exception as e:
                    manager_session.rollback()
                    error_trace = traceback.format_exc()
                    logger.error(
                        f"Unexpected error while provisioning database: {e}\nStack trace:\n{error_trace}"
                    )
                    return QueryResult(
                        success=False,
                        session_id=user_session_id,
                        message=f"An unexpected error occurred. Please try again later. {e}",
                        instance_id=instance_id,
                    )
            user_session_engine = create_user_session_engine(user_session)

        success = True
        runs: list[QueryRun] = []
        with user_session_engine.connect() as conn:
            # Set session tidb_multi_statement_mode to 'ON' to reduce errors
            try:
                conn.execute(text("SET SESSION tidb_multi_statement_mode='ON'"))
            except Exception as e:
                logger.error(f"Failed to set tidb_multi_statement_mode: {e}")

            for step in query_request.steps:
                try:
                    logger.debug(f"Executing SQL: {step.sql}")
                    start_time = time.time()
                    res = conn.execute(text(step.sql))
                    end_time = time.time()
                    execution_time = end_time - start_time
                    if res.returns_rows:
                        results = [list(row) for row in res.all()]
                    else:
                        results = [
                            [
                                f"Query OK, {res.rowcount} row affected ({execution_time:.3f} sec)"
                            ]
                        ]

                    results = serial_helper(results)
                    # Use LLM to verify the results
                    verification_result = sql_result_verifier.verify(
                        json.dumps([run.model_dump() for run in runs]),
                        step,
                        str(results),
                    )

                    runs.append(
                        QueryRun(
                            explanation=step.explanation,
                            sql=step.sql,
                            success=verification_result.excepted,
                            results=results,
                            llm_verification=verification_result.message,
                        )
                    )

                    if not verification_result.excepted:
                        success = False
                        if query_request.fast_fail:
                            break
                except sqlalchemy.exc.ProgrammingError as e:
                    logger.error(
                        f"ProgrammingError executing SQL: {step.sql}, error: {e}"
                    )
                    # Use LLM to verify if this error is expected
                    verification_result = sql_result_verifier.verify(
                        json.dumps([run.model_dump() for run in runs]), step, str(e)
                    )

                    runs.append(
                        QueryRun(
                            explanation=step.explanation,
                            sql=step.sql,
                            success=verification_result.excepted,
                            results=[[str(e)]],
                            sql_error_code=e.orig.args[0],
                            sql_error_message=e.orig.args[1],
                            llm_verification=verification_result.message,
                        )
                    )

                    if not verification_result.excepted:
                        success = False
                        if query_request.fast_fail:
                            break
                except Exception as e:
                    logger.error(
                        f"Exception error executing SQL: {step.sql}, error: {e}"
                    )
                    # Use LLM to verify if this error is expected
                    verification_result = sql_result_verifier.verify(
                        json.dumps([run.model_dump() for run in runs]), step, str(e)
                    )

                    runs.append(
                        QueryRun(
                            explanation=step.explanation,
                            sql=step.sql,
                            success=verification_result.excepted,
                            results=[[str(e)]],
                            sql_error_message=str(e),
                            llm_verification=verification_result.message,
                        )
                    )

                    if not verification_result.excepted:
                        success = False
                        if query_request.fast_fail:
                            break
        try:
            if instance_count > 20:
                with Scoped_Session() as manager_session:
                    dispose_instance(manager_session, instance)
            else:
                dispose_database_for_session(user_session)
        except Exception as e:
            logger.error(
                f"Failed to dispose database (instance count = {instance_count}) for job: {e}"
            )

        return QueryResult(
            success=success,
            session_id=user_session_id,
            runs=runs,
            instance_id=instance_id,
        )
    except Exception as e:
        logger.error(f"Internal server error: {e}")
        return QueryResult(
            success=False,
            session_id=user_session_id,
            message=f"Internal server error: {e}",
            instance_id=instance_id,
        )


def provision_database_for_session(
    manager_session: Session, instance: DBInstance, user_session: UserSession
):
    short_id = user_session.id.hex[:8]
    instance_engine = create_instance_engine(instance)

    @retry(
        stop=stop_after_attempt(4), wait=wait_exponential(multiplier=5, min=5, max=30)
    )
    def prepare_database():
        with instance_engine.connect() as conn:
            with conn.begin():
                db_name = f"db_{short_id}"
                conn.execute(text(f"CREATE DATABASE IF NOT EXISTS {db_name};"))

                if instance.tidb_cloud_branch_id is None:
                    prefix = instance.db_user.split(".")[0]
                    db_user = f"{prefix}.u_{short_id}"
                    db_password = generate_random_password()

                    conn.execute(
                        text(
                            f"CREATE USER IF NOT EXISTS '{db_user}' IDENTIFIED BY '{db_password}';"
                        )
                    )
                    conn.execute(
                        text(f"GRANT ALL PRIVILEGES ON {db_name}.* TO '{db_user}';")
                    )
                else:
                    db_user = instance.db_user
                    db_password = instance.db_password

                return db_user, db_password, db_name

    try:
        db_user, db_password, db_name = prepare_database()

        user_session.db_host = instance.db_host
        user_session.db_port = instance.db_port
        user_session.db_user = db_user
        user_session.db_password = db_password
        user_session.db_name = db_name
        user_session.instance_id = instance.id

        manager_session.add(user_session)
        manager_session.commit()

        return user_session
    except SQLAlchemyError as e:
        logger.error(f"Failed to provision database for user session: {e}")
        manager_session.rollback()
        raise
    except Exception as e:
        logger.error(f"Unexpected error while provisioning database: {e}")
        manager_session.rollback()
        raise


def dispose_database_for_session(user_session: UserSession):
    instance_engine = create_user_session_engine(user_session)
    with instance_engine.connect() as conn:
        # drop the database schema for user session.
        conn.execute(text(f"DROP DATABASE IF EXISTS {user_session.db_name};"))


def dispose_instance(manager_session: Session, instance: DBInstance):
    db_instance_repo.update_instance_status(
        manager_session, instance, InstanceStatus.INACTIVE
    )


# Generate a random password
def generate_random_password(length=12):
    characters = string.ascii_uppercase + string.digits
    return "".join(secrets.choice(characters) for _ in range(length))


class NotDBInstanceAvailable(NoResultFound):
    pass


class FailedToApplyDBInstance(Exception):
    pass


def create_instance_engine(instance: DBInstance):
    cluster_name = (
        f"instance_engine:{instance.tidb_cloud_cluster_name}" or "instance_engine"
    )

    logger.info(f"Creating engine for cluster: {cluster_name}")
    database_url = str(
        MultiHostUrl.build(
            scheme="mysql+pymysql",
            username=quote_plus(instance.db_user),
            password=quote_plus(instance.db_password),
            host=instance.db_host,
            port=instance.db_port,
            query=(
                "ssl_verify_cert=true&ssl_verify_identity=true"
                if settings.TIDB_SSL
                else None
            ),
        )
    )
    return create_engine(database_url, echo=False, logging_name=cluster_name)


def create_user_session_engine(session: UserSession):
    engine_name = f"session_engine:{session.db_user}" or "session_engine"
    database_url = str(
        MultiHostUrl.build(
            scheme="mysql+pymysql",
            username=quote_plus(session.db_user),
            password=quote_plus(session.db_password),
            host=session.db_host,
            port=int(session.db_port),
            path=session.db_name,
            query=(
                "ssl_verify_cert=true&ssl_verify_identity=true"
                if settings.TIDB_SSL
                else None
            ),
        )
    )
    return create_engine(database_url, echo=True, logging_name=engine_name)


def serial_helper(data):
    """Helper function to serialize MySQL data types to JSON-compatible formats."""
    if isinstance(data, list):
        return [serial_helper(item) for item in data]
    elif isinstance(data, dict):
        return {key: serial_helper(value) for key, value in data.items()}
    elif isinstance(data, decimal.Decimal):
        # Convert Decimal to string to maintain precision
        return str(data)
    elif isinstance(data, (datetime.datetime, datetime.date, datetime.time)):
        # Convert date/time to ISO 8601 format string
        return data.isoformat()
    elif isinstance(data, bytes):
        try:
            return data.decode("utf-8")
        except UnicodeDecodeError:
            return data.hex()
    else:
        return data


def render_runs_to_markdown(runs: list[QueryRun]):
    markdown = ""

    for run in runs:
        # Add explanation as a section header
        markdown += f"#### {run.explanation}\n\n"

        # Check success and add appropriate status block (success or danger)
        if run.success:
            markdown += ":::success\n" + run.llm_verification + "\n:::\n\n"
        else:
            markdown += ":::danger\n" + run.llm_verification + "\n:::\n\n"

        # Add the SQL query in a code block
        markdown += "```sql\n"
        markdown += run.sql + "\n"
        markdown += "```\n\n"

        # Add the results or error message if present
        markdown += (
            "```\n"
            + (str(run.results) if run.results is not None else "No results")
            + "\n```\n\n"
        )

    return markdown