package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
)

type DB struct {
	pool *pgxpool.Pool
}

func initDB(ctx context.Context) (*DB, error) {
	dsn := os.Getenv("DATABASE_URL")
	if dsn == "" {
		return nil, fmt.Errorf("DATABASE_URL environment variable not set")
	}
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		return nil, err
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, err
	}
	// Run migrations on startup
	if err := runMigrations(ctx, pool); err != nil {
		log.Printf("WARNING: Migrations failed (will continue anyway): %v", err)
	} else {
		log.Println("PostgreSQL connected, all migrations applied")
	}

	db := &DB{pool: pool}
	// Ensure SYS-014 special layers exist
	db.ensureSYS014Layers(ctx)

	return db, nil
}

func (db *DB) Close() { db.pool.Close() }

// ensureSYS014Layers ensures SYS-014 special layers exist in the database
func (db *DB) ensureSYS014Layers(ctx context.Context) {
	// First, ensure constraint allows high layer numbers
	_, _ = db.pool.Exec(ctx, `
		ALTER TABLE ic3_nav_layer DROP CONSTRAINT IF EXISTS ic3_nav_layer_layer_no_check;
		ALTER TABLE ic3_nav_layer ADD CONSTRAINT ic3_nav_layer_layer_no_check CHECK (layer_no BETWEEN 1 AND 999);
	`)

	// Insert SYS-014 special layers
	_, _ = db.pool.Exec(ctx, `
		INSERT INTO ic3_nav_layer (layer_no, code, label, description) VALUES
		(301, 'SYS014-L1', 'Systems & Architecture', 'The 13 water systems as integrated pillars'),
		(302, 'SYS014-L2', 'Cross-System Analytics', 'Unified intelligence across all domains')
		ON CONFLICT DO NOTHING;
	`)

	// Ensure SYS-014 exists in system master
	_, _ = db.pool.Exec(ctx, `
		INSERT INTO ic3_system_master (system_id, system_name, system_label, created_at)
		VALUES ('SYS-014', 'IC3 Platform Master', 'IC³ Platform Master', NOW())
		ON CONFLICT DO NOTHING;
	`)

	log.Println("[INIT] ✓ SYS-014 layers initialized")
}

// seedDefaultUsers creates default admin user if none exist
func (db *DB) seedDefaultUsers(ctx context.Context) {
	// Check if admin user exists
	var exists bool
	err := db.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM ic3_user_master WHERE username='admin')`).Scan(&exists)
	if err != nil {
		log.Printf("seedDefaultUsers: cannot check if admin exists: %v", err)
		return
	}
	if exists {
		return
	}

	// Insert admin user with bcrypt hash of "admin123"
	_, err = db.pool.Exec(ctx, `
		INSERT INTO ic3_user_master (user_id, username, password_hash, full_name, email, is_active, created_at, updated_at)
		VALUES ('USR-001', 'admin', '$2a$10$N9qo8uLOickgx2ZMRZoMyeIjZAgcg7b3XeKeUxWdeS86E36gBS43O', 'Administrator', 'admin@ic3.local', TRUE, NOW(), NOW())
	`)
	if err != nil {
		log.Printf("seedDefaultUsers: failed to insert admin: %v", err)
		return
	}
	log.Println("✓ Seeded default admin user (password: admin123)")
}

// seedDomainSystems truncates ic3_system_master and seeds domain systems
func (db *DB) seedDomainSystems(ctx context.Context) {
	domains := []struct{ id, name string }{
		{"SYS-01", "24/7"},
		{"SYS-02", "Realtime"},
		{"SYS-03", "Water Quality"},
		{"SYS-04", "WTP"},
		{"SYS-05", "WWTP"},
		{"SYS-06", "STP"},
		{"SYS-07", "GIS"},
		{"SYS-08", "Digital Twin"},
		{"SYS-09", "Irrigation & Agri Water"},
		{"SYS-10", "RWS"},
		{"SYS-11", "GW & SW"},
		{"SYS-12", "NRW"},
	}

	// Delete all existing systems
	_, err := db.pool.Exec(ctx, `DELETE FROM ic3_system_master`)
	if err != nil {
		log.Printf("Failed to delete systems: %v", err)
		return
	}

	// Insert domain systems
	for _, d := range domains {
		_, err := db.pool.Exec(ctx,
			`INSERT INTO ic3_system_master (system_id, system_name, system_category, system_type, is_active, created_at, updated_at)
			 VALUES ($1, $2, 'DOMAIN', 'DOMAIN', TRUE, NOW(), NOW())`,
			d.id, d.name,
		)
		if err != nil {
			log.Printf("Failed to seed domain %s: %v", d.name, err)
			continue
		}
		log.Printf("Seeded domain: %s", d.name)
	}
}

// saveBatch writes each envelope to:
//   - telemetry       (append-only history)
//   - asset_latest    (upsert – latest state per asset)
//   - assets          (upsert – registry, tracks first/last seen)
//   - alarm_events    (insert on alarm transition away from NORMAL)
func (db *DB) saveBatch(ctx context.Context, batch []TelemetryEnvelope) {
	for _, env := range batch {
		payload, _ := json.Marshal(env.Values)

		var ts any
		if env.OccurredAt != "" {
			ts = env.OccurredAt
		} else {
			ts = time.Now().UTC()
		}

		// 1. Append to telemetry history
		db.pool.Exec(ctx, `
			INSERT INTO telemetry (
				record_id, timestamp_utc, domain_code, system_id, asset_id,
				source_system, collector_id, sequence_no,
				quality_code, alarm_state,
				location_path, country_code, state_code, district_code,
				taluk_code, city_code, ward_code,
				zone_id, dma_id, pressure_zone_id, site_id,
				latitude, longitude, elevation_m,
				stream_topic, priority, payload
			) VALUES (
				$1,$2,$3,$4,$5,
				$6,$7,$8,
				$9,$10,
				$11,$12,$13,$14,
				$15,$16,$17,
				$18,$19,$20,$21,
				$22,$23,$24,
				$25,$26,$27
			)
			ON CONFLICT (record_id) DO NOTHING`,
			nullText(env.RecordID), ts, env.Domain, env.SystemID, env.AssetID,
			nullText(env.SourceSystem), nullText(env.CollectorID), nullInt(env.SequenceNo),
			coalesce(env.Quality, "GOOD"), coalesce(env.Alarm, "NORMAL"),
			nullText(env.LocationPath), nullText(env.CountryCode), nullText(env.StateCode), nullText(env.DistrictCode),
			nullText(env.TalukCode), nullText(env.CityCode), nullText(env.WardCode),
			nullText(env.ZoneID), nullText(env.DMAID), nullText(env.PressureZoneID), nullText(env.SiteID),
			nullFloat(env.Latitude), nullFloat(env.Longitude), nullFloat(env.ElevationM),
			nullText(env.StreamTopic), nullText(env.Priority), payload,
		)

		// 2. Upsert asset_latest
		db.pool.Exec(ctx, `
			INSERT INTO asset_latest (
				asset_id, system_id, domain_code, timestamp_utc,
				quality_code, alarm_state, dma_id, site_id, stream_topic, payload
			) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10)
			ON CONFLICT (asset_id) DO UPDATE SET
				system_id     = EXCLUDED.system_id,
				domain_code   = EXCLUDED.domain_code,
				timestamp_utc = EXCLUDED.timestamp_utc,
				quality_code  = EXCLUDED.quality_code,
				alarm_state   = EXCLUDED.alarm_state,
				dma_id        = EXCLUDED.dma_id,
				site_id       = EXCLUDED.site_id,
				stream_topic  = EXCLUDED.stream_topic,
				payload       = EXCLUDED.payload,
				updated_at    = NOW()`,
			env.AssetID, env.SystemID, env.Domain, ts,
			coalesce(env.Quality, "GOOD"), coalesce(env.Alarm, "NORMAL"),
			nullText(env.DMAID), nullText(env.SiteID), nullText(env.StreamTopic), payload,
		)

		// 3. Upsert asset registry
		db.pool.Exec(ctx, `
			INSERT INTO assets (
				asset_id, system_id, domain_code, site_id, dma_id, zone_id,
				location_path, latitude, longitude, elevation_m,
				first_seen, last_seen, last_quality, last_alarm, total_records
			) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10, $11,$11,$12,$13,1)
			ON CONFLICT (asset_id) DO UPDATE SET
				last_seen     = EXCLUDED.last_seen,
				last_quality  = EXCLUDED.last_quality,
				last_alarm    = EXCLUDED.last_alarm,
				total_records = assets.total_records + 1`,
			env.AssetID, env.SystemID, env.Domain, nullText(env.SiteID), nullText(env.DMAID), nullText(env.ZoneID),
			nullText(env.LocationPath), nullFloat(env.Latitude), nullFloat(env.Longitude), nullFloat(env.ElevationM),
			ts, coalesce(env.Quality, "GOOD"), coalesce(env.Alarm, "NORMAL"),
		)

		// 4. Record alarm transitions (NORMAL→anything or anything→NORMAL)
		if env.Alarm != "" && env.Alarm != "NORMAL" {
			alarmID := env.RecordID + ":alarm"
			if alarmID == ":alarm" {
				alarmID = env.AssetID + ":" + fmt.Sprint(time.Now().UnixNano())
			}
			db.pool.Exec(ctx, `
				INSERT INTO alarm_events (
					record_id, timestamp_utc, asset_id, domain_code, system_id,
					dma_id, from_state, to_state, quality_code
				) VALUES ($1,$2,$3,$4,$5,$6,'NORMAL',$7,$8)
				ON CONFLICT (record_id) DO NOTHING`,
				alarmID, ts, env.AssetID, env.Domain, env.SystemID,
				nullText(env.DMAID), env.Alarm, coalesce(env.Quality, "GOOD"),
			)
		}
	}
}

// history queries telemetry with the v3 column names.
func (db *DB) history(ctx context.Context, assetID, domain string, limit int) ([]TelemetryEnvelope, error) {
	if db == nil {
		return nil, nil
	}
	q := `SELECT asset_id, domain_code, system_id, timestamp_utc,
	             quality_code, alarm_state, payload
	      FROM telemetry WHERE 1=1`
	args := []any{}
	n := 1
	if assetID != "" {
		q += fmt.Sprintf(" AND asset_id=$%d", n)
		args = append(args, assetID)
		n++
	}
	if domain != "" {
		q += fmt.Sprintf(" AND domain_code=$%d", n)
		args = append(args, domain)
		n++
	}
	q += fmt.Sprintf(" ORDER BY created_at DESC LIMIT $%d", n)
	args = append(args, limit)

	rows, err := db.pool.Query(ctx, q, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var out []TelemetryEnvelope
	for rows.Next() {
		var e TelemetryEnvelope
		var raw []byte
		var ts *time.Time
		if err := rows.Scan(&e.AssetID, &e.Domain, &e.SystemID, &ts, &e.Quality, &e.Alarm, &raw); err != nil {
			continue
		}
		if raw != nil {
			json.Unmarshal(raw, &e.Values)
		}
		if ts != nil {
			e.OccurredAt = ts.UTC().Format(time.RFC3339)
		}
		out = append(out, e)
	}
	return out, nil
}

// latestAssets reads the current state of every asset from asset_latest.
// Optionally filtered by domain. Returns TelemetryEnvelope slice so existing
// callers and the WS broadcast format stay unchanged.
func (db *DB) latestAssets(ctx context.Context, domain string) ([]TelemetryEnvelope, error) {
	if db == nil {
		return nil, nil
	}
	q := `SELECT asset_id, system_id, domain_code, timestamp_utc,
	             quality_code, alarm_state, dma_id, site_id, stream_topic, payload, updated_at
	      FROM asset_latest`
	args := []any{}
	if domain != "" {
		q += ` WHERE domain_code = $1`
		args = append(args, domain)
	}
	q += ` ORDER BY updated_at DESC`

	rows, err := db.pool.Query(ctx, q, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var out []TelemetryEnvelope
	for rows.Next() {
		var e TelemetryEnvelope
		var raw []byte
		var ts *time.Time
		var updatedAt time.Time
		if err := rows.Scan(
			&e.AssetID, &e.SystemID, &e.Domain, &ts,
			&e.Quality, &e.Alarm, &e.DMAID, &e.SiteID, &e.StreamTopic, &raw, &updatedAt,
		); err != nil {
			continue
		}
		if ts != nil {
			e.OccurredAt = ts.UTC().Format(time.RFC3339)
		}
		if raw != nil {
			json.Unmarshal(raw, &e.Values)
		}
		out = append(out, e)
	}
	return out, nil
}

// latestAlarms reads only assets currently in a non-NORMAL alarm state.
func (db *DB) latestAlarms(ctx context.Context) ([]map[string]any, error) {
	if db == nil {
		return []map[string]any{}, nil
	}
	rows, err := db.pool.Query(ctx, `
		SELECT asset_id, domain_code, system_id, timestamp_utc,
		       quality_code, alarm_state, dma_id, site_id, payload
		FROM asset_latest
		WHERE alarm_state <> 'NORMAL'
		ORDER BY updated_at DESC`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var out []map[string]any
	for rows.Next() {
		var assetID, domain, sysID, quality, alarm string
		var dmaID, siteID *string
		var ts *time.Time
		var raw []byte
		if err := rows.Scan(&assetID, &domain, &sysID, &ts, &quality, &alarm, &dmaID, &siteID, &raw); err != nil {
			continue
		}
		var values map[string]any
		if raw != nil {
			json.Unmarshal(raw, &values)
		}
		var tsStr string
		if ts != nil {
			tsStr = ts.UTC().Format(time.RFC3339)
		}
		out = append(out, map[string]any{
			"asset_id":    assetID,
			"alarm_state": alarm,
			"quality":     quality,
			"domain":      domain,
			"system_id":   sysID,
			"dma_id":      dmaID,
			"site_id":     siteID,
			"ts":          tsStr,
			"values":      values,
		})
	}
	if out == nil {
		out = []map[string]any{}
	}
	return out, nil
}

// stats returns aggregate counts from the DB tables.
func (db *DB) stats(ctx context.Context) (map[string]any, error) {
	if db == nil {
		return map[string]any{"total": 0, "alarms": 0, "domains": 0, "live_assets": 0}, nil
	}
	var total, domains int
	var liveAssets, liveAlarms int
	db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM telemetry`).Scan(&total)
	db.pool.QueryRow(ctx, `SELECT COUNT(DISTINCT domain_code) FROM asset_latest`).Scan(&domains)
	db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM asset_latest`).Scan(&liveAssets)
	db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM asset_latest WHERE alarm_state <> 'NORMAL'`).Scan(&liveAlarms)
	return map[string]any{
		"total":       total,
		"alarms":      liveAlarms,
		"domains":     domains,
		"live_assets": liveAssets,
	}, nil
}

// statsV2 returns aggregate counts from both legacy and v5.0 tables
func (db *DB) statsV2(ctx context.Context) (map[string]any, error) {
	if db == nil {
		return map[string]any{
			"total": 0, "alarms": 0, "domains": 0, "live_assets": 0,
			"ai_anomalies": 0, "pump_critical": 0, "nrw_breach_count": 0, "predictive_wo": 0,
		}, nil
	}

	// Existing queries
	var total, domains int
	var liveAssets, liveAlarms int
	db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM telemetry`).Scan(&total)
	db.pool.QueryRow(ctx, `SELECT COUNT(DISTINCT domain_code) FROM asset_latest`).Scan(&domains)
	db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM asset_latest`).Scan(&liveAssets)
	db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM asset_latest WHERE alarm_state <> 'NORMAL'`).Scan(&liveAlarms)

	// New v5.0 queries
	var aiAnomalies, pumpCritical, nrwBreach, predictiveWO int
	db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM ai_anomaly_detection WHERE created_at > NOW()-INTERVAL'1 hour'`).Scan(&aiAnomalies)
	db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM ai_asset_health_score WHERE risk_level='CRITICAL' AND calculated_time > NOW()-INTERVAL'2 hours'`).Scan(&pumpCritical)
	db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM nrw_dma_water_balance WHERE nrw_percentage > 20 AND calculation_date = CURRENT_DATE`).Scan(&nrwBreach)
	db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM cmms_work_order WHERE wo_type='PREDICTIVE' AND status='OPEN'`).Scan(&predictiveWO)

	return map[string]any{
		"total":            total,
		"alarms":           liveAlarms,
		"domains":          domains,
		"live_assets":      liveAssets,
		"ai_anomalies":     aiAnomalies,
		"pump_critical":    pumpCritical,
		"nrw_breach_count": nrwBreach,
		"predictive_wo":    predictiveWO,
	}, nil
}

// ── helpers ───────────────────────────────────────────────────────────────────

func coalesce(s, def string) string {
	if s == "" {
		return def
	}
	return s
}

func nullText(s string) any {
	if s == "" {
		return nil
	}
	return s
}

func nullFloat(f float64) any {
	if f == 0 {
		return nil
	}
	return f
}

func nullInt(i int) any {
	if i == 0 {
		return nil
	}
	return i
}

// ============================================================================
// SYSTEM MENU MANAGEMENT
// ============================================================================

// ListMenuItems retrieves all menu items
func (db *DB) ListMenuItems(ctx context.Context) ([]MenuItem, error) {
	rows, err := db.pool.Query(ctx, `
		SELECT id, label, domain, role, visible, description, "order", created_at, updated_at
		FROM ic3_system_menu
		ORDER BY domain, "order" ASC
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var items []MenuItem
	for rows.Next() {
		var item MenuItem
		err := rows.Scan(&item.ID, &item.Label, &item.Domain, &item.Role, &item.Visible,
			&item.Description, &item.Order, &item.CreatedAt, &item.UpdatedAt)
		if err != nil {
			return nil, err
		}
		items = append(items, item)
	}
	return items, rows.Err()
}

// GetMenuItem retrieves a single menu item by ID
func (db *DB) GetMenuItem(ctx context.Context, id int) (*MenuItem, error) {
	var item MenuItem
	err := db.pool.QueryRow(ctx, `
		SELECT id, label, domain, role, visible, description, "order", created_at, updated_at
		FROM ic3_system_menu
		WHERE id = $1
	`, id).Scan(&item.ID, &item.Label, &item.Domain, &item.Role, &item.Visible,
		&item.Description, &item.Order, &item.CreatedAt, &item.UpdatedAt)
	if err != nil {
		return nil, err
	}
	return &item, nil
}

// CreateMenuItem inserts a new menu item
func (db *DB) CreateMenuItem(ctx context.Context, item *MenuItem) (*MenuItem, error) {
	var id int
	var createdAt, updatedAt string

	err := db.pool.QueryRow(ctx, `
		INSERT INTO ic3_system_menu (label, domain, role, visible, description, "order", created_at, updated_at)
		VALUES ($1, $2, $3, $4, $5, $6, NOW(), NOW())
		RETURNING id, created_at, updated_at
	`, item.Label, item.Domain, item.Role, item.Visible, item.Description, item.Order).
		Scan(&id, &createdAt, &updatedAt)

	if err != nil {
		return nil, err
	}

	return &MenuItem{
		ID:          id,
		Label:       item.Label,
		Domain:      item.Domain,
		Role:        item.Role,
		Visible:     item.Visible,
		Description: item.Description,
		Order:       item.Order,
		CreatedAt:   createdAt,
		UpdatedAt:   updatedAt,
	}, nil
}

// UpdateMenuItem updates an existing menu item
func (db *DB) UpdateMenuItem(ctx context.Context, item *MenuItem) (*MenuItem, error) {
	var updatedAt string

	err := db.pool.QueryRow(ctx, `
		UPDATE ic3_system_menu
		SET label = $1, domain = $2, role = $3, visible = $4, description = $5, "order" = $6, updated_at = NOW()
		WHERE id = $7
		RETURNING updated_at
	`, item.Label, item.Domain, item.Role, item.Visible, item.Description, item.Order, item.ID).
		Scan(&updatedAt)

	if err != nil {
		return nil, err
	}

	item.UpdatedAt = updatedAt
	return item, nil
}

// DeleteMenuItem removes a menu item
func (db *DB) DeleteMenuItem(ctx context.Context, id int) error {
	result, err := db.pool.Exec(ctx, `DELETE FROM ic3_system_menu WHERE id = $1`, id)
	if err != nil {
		return err
	}
	if result.RowsAffected() == 0 {
		return fmt.Errorf("menu item not found")
	}
	return nil
}