import json
import datetime
from sqlalchemy import (
    Column,
    String,
    Integer,
    Text,
    DateTime,
    BigInteger,
)
from sqlalchemy.dialects.mysql import LONGTEXT

from consts.consts import EdaTaskJobStatus
from models.base import Base, BaseMixin, with_session


class EDATask(Base, BaseMixin):
    __tablename__ = "eda_task"

    name = Column(String(128), nullable=False)
    url_path = Column(String(128), nullable=False, unique=True)
    method = Column(String(8), nullable=False)
    description = Column(String(256), nullable=False)
    timeout = Column(Integer, nullable=False)
    prompt_template = Column(Text, nullable=False)
    exprs = Column(String(1024), nullable=False)

    @classmethod
    def create(cls, s, name, url_path, method, description, timeout, prompt_template, exprs):
        item = cls(
            name=name, url_path=url_path, method=method, description=description, timeout=timeout,
            prompt_template=prompt_template, exprs=json.dumps(exprs),
        )
        s.add(item)
        s.flush()

        return item

    @classmethod
    def get_by_url_path(cls, s, url_path):
        return s.query(cls).filter(cls.url_path == url_path).first()

    def update(self, s, name, url_path, method, description, timeout, prompt_template, exprs):
        self.name = name
        self.url_path = url_path
        self.method = method
        self.description = description
        self.timeout = timeout
        self.prompt_template = prompt_template
        self.exprs = json.dumps(exprs)

        s.flush()

    def load_exprs(self):
        return json.loads(self.exprs)

    def to_json(self):
        return {
            "id": self.id,
            "name": self.name,
            "url_path": self.url_path,
            "method": self.method,
            "description": self.description,
            "timeout": self.timeout,
            "prompt_template": self.prompt_template,
            "exprs": json.loads(self.exprs) if self.exprs else None,
        }


class EDATaskJob(Base, BaseMixin):
    __tablename__ = "eda_task_job"

    eda_task_id = Column(Integer, nullable=False)
    job_id = Column(String(36), nullable=False, index=True)
    status = Column(String(8), nullable=False, default=EdaTaskJobStatus.INIT)
    prompt = Column(Text, nullable=False, default="")
    output = Column(Text, nullable=False, default="")
    result = Column(LONGTEXT, nullable=False, default="")
    ended_at = Column(DateTime)
    variables = Column(String(1024), nullable=False)
    fn_name = Column(String(128))
    fn_args = Column(String(4096))
    last_updated = Column(DateTime, default=datetime.datetime.now)
    org_id = Column(String(64))

    @classmethod
    def create(cls, s, org_id, job_id, prompt, variables, fn_name=None, fn_args=None):
        eda_task_id = 1  # magic number for history reason
        item = cls(
            org_id=org_id, eda_task_id=eda_task_id, job_id=job_id, prompt=prompt, variables=json.dumps(variables),
            fn_name=fn_name, fn_args=json.dumps(fn_args),
        )
        s.add(item)
        s.flush()

        return item

    @classmethod
    def get_all_running_jobs_before(cls, s, before, for_update=False):
        q = s.query(cls).filter(cls.status == EdaTaskJobStatus.RUNNING, cls.updated_at < before)

        if for_update:
            q = q.with_for_update()

        return q.all()

    @classmethod
    def get_all_running_jobs_after(cls, s, after):
        return s.query(cls).filter(cls.status == EdaTaskJobStatus.RUNNING, cls.updated_at > after).all()

    def refresh_last_updated(self, s):
        self.last_updated = datetime.datetime.now()
        s.flush()

    @classmethod
    def get_by_job_id(cls, s, job_id):
        return s.query(cls).filter(cls.job_id == job_id).first()

    def set_result(self, s, result):
        self.result = json.dumps(result, default=str)
        self.refresh_last_updated(s)
        s.flush()

    def load_result(self):
        return json.loads(self.result) if self.result else None

    def append_output(self, s, output):
        self.output += output
        s.flush()

    def update_status(self, s, status):
        self.status = status
        if self.status in (EdaTaskJobStatus.DONE, EdaTaskJobStatus.FAILED):
            self.ended_at = datetime.datetime.now()

        s.flush()

    def load_variables(self):
        return json.loads(self.variables)

    def to_json(self):
        ended_at = 0
        if self.ended_at:
            ended_at = int(self.ended_at.timestamp())

        result = self.load_result()
        if isinstance(result, dict):
            if result:
                result.pop("aipower_resp", None)

        job_detail = {
            "job_id": self.job_id,
            "status": self.status,
            "output": self.output,
            "result": result,
            "ended_at": ended_at,
            "variables": self.load_variables(),
        }

        reason = ""
        if self.status == EdaTaskJobStatus.FAILED:
            reason = ""
            if job_detail.get("result"):
                reason = job_detail["result"].get("reason")

            job_detail["result"] = None

        job_detail["reason"] = reason

        return job_detail