"""
database.py — SQLite persistence via SQLAlchemy.

Tables: runs, results, schedules + compliance sprint tables
(baseline_profiles, delta_reports, vulnerability_deltas,
 compliance_audit_trails, mitigations, sign_offs, standards_mappings)
"""
from __future__ import annotations

import json
from datetime import datetime
from pathlib import Path
from typing import Optional

from sqlalchemy import (
    Boolean, Column, DateTime, Float, Integer, String, Text,
    create_engine, event, text,
)
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker

_DB_PATH = Path.home() / ".sichgate" / "sichgate.db"


class Base(DeclarativeBase):
    pass


class RunRecord(Base):
    __tablename__ = "runs"

    id = Column(String(36), primary_key=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    model_target = Column(String(512))
    attacks_requested = Column(Text)   # JSON array of attack names
    status = Column(String(20))        # "completed" | "running" | "failed"
    summary_json = Column(Text)        # Full RunSummary JSON


class ResultRecord(Base):
    __tablename__ = "results"

    id = Column(String(36), primary_key=True)
    run_id = Column(String(36))
    attack_name = Column(String(128))
    attack_type = Column(String(64))
    status = Column(String(20))
    severity = Column(String(20))
    full_json = Column(Text)           # Complete AttackResult JSON


class ScheduleRecord(Base):
    __tablename__ = "schedules"

    id = Column(String(36), primary_key=True)
    model_target = Column(String(512))
    attacks = Column(Text)             # JSON array
    cron_expression = Column(String(64))
    enabled = Column(Boolean, default=True)
    last_run_at = Column(DateTime, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)


# ---------------------------------------------------------------------------
# Gap #1: Comparative Vulnerability Profiling tables
# ---------------------------------------------------------------------------

class BaselineProfileRecord(Base):
    __tablename__ = "baseline_profiles"

    id = Column(String(36), primary_key=True)          # UUID
    model_id = Column(String(512))                     # HF model ID or path
    model_hash = Column(String(64), nullable=True)     # SHA256 of weights
    model_version = Column(String(128), nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    attacks_run = Column(Text)                         # JSON array of attack names
    results_json = Column(Text)                        # JSON list of AttackResult dicts
    param_count_millions = Column(Integer, nullable=True)


class DeltaReportRecord(Base):
    __tablename__ = "delta_reports"

    id = Column(String(36), primary_key=True)          # UUID
    baseline_id = Column(String(36))                   # FK → baseline_profiles.id
    ft_model_id = Column(String(512))                  # Fine-tuned model path
    ft_model_hash = Column(String(64), nullable=True)
    ft_data_path = Column(String(1024), nullable=True) # Training data path
    created_at = Column(DateTime, default=datetime.utcnow)
    summary_json = Column(Text)                        # Full DeltaReport JSON
    fixed_count = Column(Integer, default=0)
    regressed_count = Column(Integer, default=0)
    neutral_count = Column(Integer, default=0)


class VulnerabilityDeltaRecord(Base):
    __tablename__ = "vulnerability_deltas"

    id = Column(String(36), primary_key=True)
    delta_report_id = Column(String(36))               # FK → delta_reports.id
    attack_type = Column(String(128))
    attack_name = Column(String(128))
    base_success_rate = Column(Float, default=0.0)
    ft_success_rate = Column(Float, default=0.0)
    classification = Column(String(20))                # FIXED | REGRESSED | NEUTRAL
    confidence = Column(Float, default=0.0)
    attribution = Column(String(256), nullable=True)   # Source analysis


# ---------------------------------------------------------------------------
# Gap #3: Regulatory Audit Trail tables
# ---------------------------------------------------------------------------

class ComplianceAuditTrailRecord(Base):
    __tablename__ = "compliance_audit_trails"

    id = Column(String(36), primary_key=True)
    finding_id = Column(String(36), unique=True)       # Links to vulnerability
    run_id = Column(String(36), nullable=True)         # Links to runs table
    attack_type = Column(String(128))
    attack_name = Column(String(128))
    seed = Column(Text, nullable=True)                 # Reproducibility seed
    payload = Column(Text, nullable=True)              # Exact prompt payload
    vulnerability_type = Column(String(128))
    status = Column(String(32), default="NEW")         # NEW|MITIGATED|ACCEPTED_RISK|RESOLVED
    created_at = Column(DateTime, default=datetime.utcnow)
    last_tested_at = Column(DateTime, nullable=True)
    standards_json = Column(Text, default="{}")        # JSON mapping of standards


class MitigationRecord(Base):
    __tablename__ = "mitigations"

    id = Column(String(36), primary_key=True)
    finding_id = Column(String(36))                    # FK → compliance_audit_trails.finding_id
    added_by = Column(String(256))
    description = Column(Text)
    mitigation_type = Column(String(64))               # content_filter|guardrails|fine-tuning|etc.
    applied_at = Column(DateTime, default=datetime.utcnow)
    before_test_json = Column(Text, nullable=True)     # AttackResult before mitigation
    after_test_json = Column(Text, nullable=True)      # AttackResult after retest
    retest_scheduled_at = Column(DateTime, nullable=True)


class SignOffRecord(Base):
    __tablename__ = "sign_offs"

    id = Column(String(36), primary_key=True)
    finding_id = Column(String(36))                    # FK → compliance_audit_trails.finding_id
    approver = Column(String(256))
    role = Column(String(64))                          # compliance_officer|ciso|legal|ai_engineer
    approved_at = Column(DateTime, default=datetime.utcnow)
    notes = Column(Text, nullable=True)


class StandardsMappingRecord(Base):
    __tablename__ = "standards_mappings"

    id = Column(String(36), primary_key=True)
    vulnerability_type = Column(String(128), unique=True)
    eu_ai_act_articles = Column(Text, default="[]")    # JSON list
    nist_ai_rmf_funcs = Column(Text, default="[]")     # JSON list
    hipaa_rules = Column(Text, default="[]")           # JSON list
    gdpr_articles = Column(Text, default="[]")         # JSON list
    owasp_refs = Column(Text, default="[]")            # JSON list


# ---------------------------------------------------------------------------
# Phase 2.1: Sovereign AIBoM, Attestation, and Differential Redteaming tables
# ---------------------------------------------------------------------------

class AIBoMRecord(Base):
    """Stores a CycloneDX 1.6 ML-BOM snapshot captured at scan time."""
    __tablename__ = "aibom_records"

    id = Column(String(36), primary_key=True)          # UUID (= test_run_id)
    run_id = Column(String(36), nullable=True)          # FK → runs.id (if associated)
    model_name = Column(String(512))
    model_version = Column(String(128), nullable=True)
    model_path = Column(String(1024))
    model_sha256 = Column(String(64), nullable=True)
    quantization_type = Column(String(64), nullable=True)
    hardware_json = Column(Text, nullable=True)         # HardwareContext JSON
    cyclonedx_json = Column(Text)                       # Full CycloneDX 1.6 BOM
    created_at = Column(DateTime, default=datetime.utcnow)


class AttestationRecord(Base):
    """Stores a cryptographic attestation signature for a report."""
    __tablename__ = "attestation_records"

    id = Column(String(36), primary_key=True)           # attestation_id UUID
    run_id = Column(String(36), nullable=True)          # FK → runs.id
    report_type = Column(String(64))                    # "findings" | "differential" | "transition"
    finding_id = Column(String(36), nullable=True)      # For status-transition attestations
    signature = Column(Text)                            # URL-safe base64 ed25519 sig
    public_key_b64 = Column(Text)                       # URL-safe base64 public key
    public_key_pem = Column(Text)                       # PEM public key
    payload_hash = Column(String(64))                   # SHA-256 of signed payload
    algorithm = Column(String(32), default="ed25519")
    signed_at = Column(DateTime, default=datetime.utcnow)


class DifferentialReportRecord21(Base):
    """Stores a Phase 2.1 differential redteaming report."""
    __tablename__ = "differential_reports_21"

    id = Column(String(36), primary_key=True)           # run_id UUID
    base_model_id = Column(String(512))
    custom_model_id = Column(String(512))
    overall_risk_delta = Column(Float, default=0.0)
    critical_alert_count = Column(Integer, default=0)
    worsened_count = Column(Integer, default=0)
    improved_count = Column(Integer, default=0)
    unchanged_count = Column(Integer, default=0)
    base_aibom_id = Column(String(36), nullable=True)   # FK → aibom_records.id
    custom_aibom_id = Column(String(36), nullable=True) # FK → aibom_records.id
    attestation_id = Column(String(36), nullable=True)  # FK → attestation_records.id
    report_json = Column(Text)                          # Full DifferentialReport JSON
    created_at = Column(DateTime, default=datetime.utcnow)


class TechnicalFileRecord(Base):
    """Stores a generated EU AI Act Article 11 Technical File."""
    __tablename__ = "technical_file_records"

    id = Column(String(36), primary_key=True)           # test_run_id UUID
    model_name = Column(String(512))
    model_hash = Column(String(64), nullable=True)
    test_run_id = Column(String(36), nullable=True)      # FK → aibom_records.id
    technical_file_json = Column(Text)                   # Full TechnicalFile JSON
    generated_timestamp = Column(String(64))
    generated_by = Column(String(64), default="sichgate-pro")
    created_at = Column(DateTime, default=datetime.utcnow)


def get_engine(db_path: Optional[Path] = None):
    path = db_path or _DB_PATH
    path.parent.mkdir(parents=True, exist_ok=True)
    engine = create_engine(f"sqlite:///{path}", echo=False)
    # WAL mode for better concurrent reads
    @event.listens_for(engine, "connect")
    def set_wal(conn, _):
        conn.execute("PRAGMA journal_mode=WAL")
    Base.metadata.create_all(engine)
    return engine


def get_session_factory(db_path: Optional[Path] = None):
    return sessionmaker(bind=get_engine(db_path))


class DatabaseManager:
    def __init__(self, db_path: Optional[Path] = None) -> None:
        self._SessionFactory = get_session_factory(db_path)

    def save_run(self, summary_dict: dict, attacks_requested: list[str]) -> None:
        with self._SessionFactory() as session:
            run = RunRecord(
                id=summary_dict["run_id"],
                model_target=summary_dict["model_target"],
                attacks_requested=json.dumps(attacks_requested),
                status="completed",
                summary_json=json.dumps(summary_dict),
            )
            session.add(run)
            for r in summary_dict.get("results", []):
                result = ResultRecord(
                    id=r.get("attack_name", "") + "_" + summary_dict["run_id"],
                    run_id=summary_dict["run_id"],
                    attack_name=r.get("attack_name", ""),
                    attack_type=r.get("attack_type", ""),
                    status=r.get("status", ""),
                    severity=r.get("severity", ""),
                    full_json=json.dumps(r),
                )
                session.add(result)
            session.commit()

    def list_runs(self, limit: int = 50, offset: int = 0) -> list[dict]:
        with self._SessionFactory() as session:
            rows = (
                session.query(RunRecord)
                .order_by(RunRecord.created_at.desc())
                .offset(offset).limit(limit).all()
            )
            return [json.loads(r.summary_json) for r in rows]

    def get_run(self, run_id: str) -> Optional[dict]:
        with self._SessionFactory() as session:
            row = session.get(RunRecord, run_id)
            return json.loads(row.summary_json) if row else None

    def save_schedule(self, schedule: dict) -> None:
        with self._SessionFactory() as session:
            rec = ScheduleRecord(
                id=schedule["id"],
                model_target=schedule["model_target"],
                attacks=json.dumps(schedule["attacks"]),
                cron_expression=schedule.get("cron_expression", "0 * * * *"),
                enabled=schedule.get("enabled", True),
            )
            session.merge(rec)
            session.commit()

    def list_schedules(self) -> list[dict]:
        with self._SessionFactory() as session:
            rows = session.query(ScheduleRecord).all()
            return [
                {
                    "id": r.id,
                    "model_target": r.model_target,
                    "attacks": json.loads(r.attacks),
                    "cron_expression": r.cron_expression,
                    "enabled": r.enabled,
                    "last_run_at": r.last_run_at.isoformat() if r.last_run_at else None,
                }
                for r in rows
            ]

    def list_baseline_profiles(self, limit: int = 50) -> list[dict]:
        with self._SessionFactory() as session:
            rows = (
                session.query(BaselineProfileRecord)
                .order_by(BaselineProfileRecord.created_at.desc())
                .limit(limit).all()
            )
            return [
                {
                    "id": r.id,
                    "model_id": r.model_id,
                    "model_hash": r.model_hash,
                    "model_version": r.model_version,
                    "created_at": r.created_at.isoformat() if r.created_at else None,
                    "attacks_run": json.loads(r.attacks_run or "[]"),
                    "param_count_millions": r.param_count_millions,
                }
                for r in rows
            ]

    def get_delta_report(self, report_id: str) -> Optional[dict]:
        with self._SessionFactory() as session:
            row = session.get(DeltaReportRecord, report_id)
            return json.loads(row.summary_json) if row else None

    def save_aibom(self, aibom_capture) -> None:
        """Persist an AIBoMCapture to the database."""
        bom = aibom_capture.to_cyclonedx()
        hw = aibom_capture.hardware.to_dict() if aibom_capture.hardware else {}
        with self._SessionFactory() as session:
            rec = AIBoMRecord(
                id=aibom_capture.test_run_id,
                model_name=aibom_capture.model_name,
                model_version=aibom_capture.model_version,
                model_path=aibom_capture.model_path,
                model_sha256=aibom_capture.model_sha256,
                quantization_type=aibom_capture.quantization_type,
                hardware_json=json.dumps(hw),
                cyclonedx_json=json.dumps(bom),
            )
            session.merge(rec)
            session.commit()

    def get_aibom(self, test_run_id: str) -> Optional[dict]:
        """Return the CycloneDX BOM dict for a given test_run_id."""
        with self._SessionFactory() as session:
            row = session.get(AIBoMRecord, test_run_id)
            return json.loads(row.cyclonedx_json) if row else None

    def save_attestation(self, attestation_block: dict, run_id: Optional[str] = None,
                         report_type: str = "findings") -> None:
        """Persist a signed attestation block from sign_findings_report()."""
        with self._SessionFactory() as session:
            rec = AttestationRecord(
                id=attestation_block["attestation_id"],
                run_id=run_id,
                report_type=report_type,
                signature=attestation_block.get("signature", ""),
                public_key_b64=attestation_block.get("public_key_b64", ""),
                public_key_pem=attestation_block.get("public_key_pem", ""),
                payload_hash=attestation_block.get("payload_hash", ""),
                algorithm=attestation_block.get("algorithm", "ed25519"),
            )
            session.merge(rec)
            session.commit()

    def get_attestation(self, attestation_id: str) -> Optional[dict]:
        """Return the attestation record dict for a given attestation_id."""
        with self._SessionFactory() as session:
            row = session.get(AttestationRecord, attestation_id)
            if not row:
                return None
            return {
                "attestation_id": row.id,
                "run_id": row.run_id,
                "report_type": row.report_type,
                "signature": row.signature,
                "public_key_b64": row.public_key_b64,
                "public_key_pem": row.public_key_pem,
                "payload_hash": row.payload_hash,
                "algorithm": row.algorithm,
                "signed_at": row.signed_at.isoformat() if row.signed_at else None,
            }

    def save_differential_report_21(self, report) -> None:
        """Persist a Phase 2.1 DifferentialReport."""
        report_dict = report.to_dict()
        analysis = report.differential_analysis
        worsened = sum(1 for d in analysis.values() if d.trend == "WORSENED")
        improved = sum(1 for d in analysis.values() if d.trend == "IMPROVED")
        unchanged = sum(1 for d in analysis.values() if d.trend == "UNCHANGED")
        with self._SessionFactory() as session:
            rec = DifferentialReportRecord21(
                id=report.run_id,
                base_model_id=report.base_model_id,
                custom_model_id=report.custom_model_id,
                overall_risk_delta=report.overall_risk_delta,
                critical_alert_count=len(report.critical_alerts),
                worsened_count=worsened,
                improved_count=improved,
                unchanged_count=unchanged,
                attestation_id=(
                    report.attestation.get("attestation_id")
                    if report.attestation else None
                ),
                report_json=json.dumps(report_dict),
            )
            session.merge(rec)
            session.commit()

    def get_differential_report_21(self, run_id: str) -> Optional[dict]:
        """Return the full Phase 2.1 differential report dict."""
        with self._SessionFactory() as session:
            row = session.get(DifferentialReportRecord21, run_id)
            return json.loads(row.report_json) if row else None

    def save_technical_file(self, technical_file) -> str:
        """Persist a TechnicalFile to the database. Returns the record ID."""
        tf_dict = technical_file.to_dict()
        model_meta = tf_dict.get("model_metadata", {})
        with self._SessionFactory() as session:
            rec = TechnicalFileRecord(
                id=technical_file.test_run_id,
                model_name=model_meta.get("name", "Unknown"),
                model_hash=model_meta.get("hash"),
                test_run_id=technical_file.test_run_id,
                technical_file_json=technical_file.to_json(),
                generated_timestamp=technical_file.timestamp,
            )
            session.merge(rec)
            session.commit()
        return technical_file.test_run_id

    def get_technical_file(self, test_run_id: str) -> Optional[dict]:
        """Return the full TechnicalFile dict for a given test_run_id."""
        with self._SessionFactory() as session:
            row = session.get(TechnicalFileRecord, test_run_id)
            return json.loads(row.technical_file_json) if row else None

    def list_technical_files(self, model_name: str = "") -> list[dict]:
        """Return metadata for all stored Technical Files, optionally filtered by model name."""
        with self._SessionFactory() as session:
            query = session.query(TechnicalFileRecord).order_by(
                TechnicalFileRecord.created_at.desc()
            )
            if model_name:
                query = query.filter(TechnicalFileRecord.model_name == model_name)
            rows = query.all()
            return [
                {
                    "id": r.id,
                    "model_name": r.model_name,
                    "model_hash": r.model_hash,
                    "test_run_id": r.test_run_id,
                    "generated_timestamp": r.generated_timestamp,
                    "generated_by": r.generated_by,
                    "created_at": r.created_at.isoformat() if r.created_at else None,
                }
                for r in rows
            ]

    def list_compliance_findings(self, limit: int = 100) -> list[dict]:
        with self._SessionFactory() as session:
            rows = (
                session.query(ComplianceAuditTrailRecord)
                .order_by(ComplianceAuditTrailRecord.created_at.desc())
                .limit(limit).all()
            )
            return [
                {
                    "id": r.id,
                    "finding_id": r.finding_id,
                    "run_id": r.run_id,
                    "attack_type": r.attack_type,
                    "attack_name": r.attack_name,
                    "vulnerability_type": r.vulnerability_type,
                    "status": r.status,
                    "created_at": r.created_at.isoformat() if r.created_at else None,
                    "standards": json.loads(r.standards_json or "{}"),
                }
                for r in rows
            ]
